Particle

競技プログラミングについての雑記

SRM 661 Div1 Medium ColorfulLineGraphs

求める値は、簡単な考察で分かる。(各 i に対して、どこに辺を張るかを考えると、辺を張るときは K-1 倍、張らないとき K 倍となる)
M が 10^6 以下と小さいので、1 周期分計算し、周期の回数(N/M回)だけかけ、周期からはみ出す分(N%M の部分)を掛け合わせれば良い。

N, K の値が大きいのでオーバーフローなどに注意

ll pow_mod(ll a, ll r, ll m ){
	ll x = 1;
	while(r){
		if(r&1) (x*=a)%=m;
		(a*=a)%=m;
		r>>=1;
	}
	return x;
}
 
class ColorfulLineGraphs {
	public:
	int countWays(long long N, long long K, int M) {
		ll mod = M;
		K--;
		if(K==0) return 1;
		ll res = 1;
		while(N%M){
			(res*=(N%mod*K%mod+1))%=mod;
			N--;
		}
		ll res2 = 1;
		rep(i, M) (res2*=i*K%mod+1)%=mod;
		(res*=pow_mod(res2, N/M, M))%=mod;
		return res;
	}
};