TopCoder SRM 661 Div1 Medium ColorfulLineGraphs

またえらい制限の差だなぁ。
http://community.topcoder.com/stat?c=problem_statement&pm=13765

問題

問題の概要はDiv2 Hardと同じ。
ただしN,Kが最大10^12と大きい。
一方、求める値はMの剰余で良い(M≦10^6)。

解法

Div2 Hardの解法は(N^(K+1))程度の計算量なので、ここでは適用できない。
TopCoder SRM 661 Div2 Hard ColorfulLineGraphsDiv2 - kmjp's blog

Div2 Hardのコードから、K=2,3の場合にNをいじるとある傾向がみられる。
すなわち、i番目の頂点及びその辺の張り方は(K+(i-1)(K-1))通りとなることがわかる。
これは以下のように考えることができる。

• i番目の点から辺を伸ばさない場合、K色の選び方はどれでも良い。
• i番目の点からj番目の点に辺を伸ばす場合、i番目はj番目と異なる色でなければならないので、i番目の色は(K-1)通り。

i番目の点から辺を伸ばす候補は(i-1)個なので、上記考察よりi番目の頂点及びその辺の張り方は(K+(i-1)(K-1))通りであることが確認できた。

あとはi=1～Nの場合を掛け合わせた(K)*(K+(K-1))*(K+2(K-1))*...*(K+(N-1)(K-1))%Mを求めればよい。
ただ、Nが非常に大きいのでこの値を愚直に計算しても間に合わない。

ここでMが小さいことに注目する。
(K)*(K+(K-1))*(K+2(K-1))*...*(K+(N-1)(K-1))の各項は等差数列であるため、周期M以下でループする。
そのようなループの部分は１周分だけ計算して、あとはループ回数の分累乗することでO(N)回の乗算を回避できる。

以下のコードでは、上記乗算を(1)*(1+(K-1))*(1+2*(K-1))...*(1+N*(K-1))と式変形し、(1+i*(K-1))==1となる周期iを検出している。

ll tot[1010101];
ll modpow(ll a, ll n,int mo) {
	ll r=1;
	while(n) r=r*((n%2)?a:1)%mo,a=a*a%mo,n>>=1;
	return r;
}

class ColorfulLineGraphs {
	public:
	int countWays(long long N, long long K, int M) {
		ll ret=1,mul=1,i;
		
		for(i=0;i<=N;i++) {
			if(i && mul==1) break;
			ret=ret*mul%M;
			tot[i]=ret;
			(mul+=K-1)%=M;
		}
		
		if(i>N) return ret;
		return modpow(ret,N/i,M)*tot[N%i]%M;
	}
}

まとめ

Div2 Hardを解いてなければこちら自力で解けなかったかも。