yukicoder : No.503 配列コレクション

最近似たようなのAtCoderで見たよな～と思いつつ、それよりは簡単なので安心した。
http://yukicoder.me/problems/no/503

問題

初期状態で全要素が1、長さがNの配列Aがある。
この配列の長さがK以上である間、長さKの部分列を選び、それらを、(部分列の要素の積×D)の1要素に置き換える、ということを繰り返す。

こうして得られる最終的な配列Bの全種類における全要素の総和を求めよ。
ただし、同一の最終系は1回ずつしかカウントしない。

解法

最終的な数列Bの要素数Mは、M=(N-1)%(K-1)+1となる。
また、部分列を選ぶ処理回数P=(N-1)/(K-1)である。

D=1の時は得られる数列は[1,1,1,...]の一択である。
以下、それ以外の場合を考える。

Bのi番目の要素B[i]は、もともと1+(K-1)*C[i]個の部分列をC[i]回置換してまとめてできたものだとする。
このとき、sum(C)=PかつB[i]=D^C[i]の形を取る。

C[i]を固定した場合、Cの残りの要素の和はP-C[i]なので、重複組み合わせよりそのような数列Bの組み合わせは通りである。
よって最終的に求める値は、C[i]を総当たりし、となる。
(先頭のMは、BのM要素について同様に考えるため)

ll N,K,D;
ll mo=1000000007;

const int NUM_=2000001;
static ll fact[NUM_+1],factr[NUM_+1],inv[NUM_+1];
ll combi(ll N_, ll C_) {
	if (fact[0]==0) {
		inv[1]=fact[0]=factr[0]=1;
		for (int i=2;i<=NUM_;++i) inv[i] = inv[mo % i] * (mo - mo / i) % mo;
		for (int i=1;i<=NUM_;++i) fact[i]=fact[i-1]*i%mo, factr[i]=factr[i-1]*inv[i]%mo;
	}
	if(C_<0 || C_>N_) return 0;
	return factr[C_]*fact[N_]%mo*factr[N_-C_]%mo;
}
ll hcomb(int P_,int Q_) { return (P_==0&&Q_==0)?1:combi(P_+Q_-1,Q_);}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>K>>D;
	
	int num=0;
	int M=(N-1)%(K-1)+1;
	int P=(N-1)/(K-1);
	
	if(D==1) {
		cout<<M<<endl;
		return;
	}
	
	ll ret=0;
	ll p=1;
	for(i=0;i<=P;i++) {
		ret += hcomb(M-1,P-i) * p %mo;
		(p*=D)%=mo;
	}
	cout<<ret*M%mo<<endl;
	
}

まとめ

500問おめでとうございます。