Codeforces #861 : Div2 E3. Minibuses on Venus (hard version)

不参加だった回。
https://codeforces.com/contest/1808/problem/E3

問題

整数N,Kが与えられる。
K進数でN桁の整数を考える。
これがラッキーであるとは、1桁を除く全桁の和と、除いた1桁の和がmod Kで一致する物を意味する。

ラッキーな整数は何通りか。

解法

• N=1またはK=1の時は0一択
• N=2の場合、ゾロ目のK通り
• K=2の時、1の個数が偶数ならよいので2^(N-1)通り
• 以下N,Kが3以上の場合を考える。
  • 桁和がSの場合、そのうち1桁をxとするとx = S-x (mod K)であるxが欲しいので2x = S (mod K)とする。
  • Kの偶奇で場合分けして考え、xの取りえる値を総当たりしよう。とはいえ各xにおいて、それを満たすN桁の整数値のバリエーションの個数は同じなので、1パターンだけ考えればよい。

ll N,K,mo;

ll modpow(ll a, ll n = mo-2) {
	ll r=1;a%=mo;
	while(n) r=r*((n%2)?a:1)%mo,a=a*a%mo,n>>=1;
	return r;
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>K>>mo;
	
	if(N==1||K==1) {
		cout<<1<<endl;
		return;
	}
	if(N==2) {
		cout<<K<<endl;
		return;
	}
	if(K==2) {
		cout<<modpow(2,N-1)<<endl;
		return;
	}
	
	ll ret=0;
	if(K%2) {
		ret+=modpow(K,N)-modpow(K-1,N)-modpow(K-1,N-1);
		x=__gcd(N-2,K);
		ret+=x*((modpow(K-1,N-1)+(K-1)*modpow(mo-1,N-1)%mo)*modpow(K%mo)%mo);
		ret+=(K-x)*((modpow(K-1,N-1)-modpow(mo-1,N-1))*modpow(K%mo)%mo);
	}
	else {
		ret+=(modpow(K,N)-modpow(K-2,N))*modpow(2)%mo-modpow(K-2,N-1);
		x=2*__gcd(N-2,K/2);
		ret+=x*(((K/2-1)*modpow(mo-2,N-1)%mo+modpow(K-2,N-1)%mo)*modpow(K%mo)%mo);
		ret+=(K-x)*((modpow(K-2,N-1)-modpow(mo-2,N-1))*modpow(K%mo)%mo);
	}
	cout<<((ret%mo+mo)%mo)<<endl;
	
}

まとめ

Codeforcesで3段階の部分点は珍しいな。