kmjp's blog

競技プログラミング参加記です

Codeforces #230 Div1 C. Yet Another Number Sequence

Codeforces

基本的なアイデアは本番中に思いついたものの、最後まで解ききれなかった。
http://codeforces.com/contest/392/problem/C

問題

F[1]=1、F[2]=2、F[i]=F[i-1]+F[i-2]で定義されるフィボナッチ数列を考える。
N,Kが与えられるので \sum_{i=1}^N (F_i \times i^K)を答えよ。

解法

かなり大きな添え字のフィボナッチ数列、ときたら行列の累乗で解くことがすぐ思い浮かぶ。
問題は行列の値である。

F_{i+1} \times (i+1)^kを漸化式で表現すると、
F_{i+1} \times (i+1)^k = (F_i + F_{i-1}) \times \( \sum_{j=1}^k {}_k C_j \times i^j \)= F_i \times \(\sum_{j=1}^k {}_k C_j \times i^j\) + F_{i-1} \times \(\sum_{j=1}^k {}_k C_j \times (i-1+1)^j\)
= F_i \times \(\sum_{j=1}^k {}_k C_j \times i^j\) + F_{i-1} \times \(\sum_{j=1}^k {}_k C_j \sum_{l=1}^j {}_j C_l \times (i-1)^j\)

よって F_{i+1} \times (i+1)^kは F_{i} \times (i)^jおよび F_{i-1} \times (i-1)^jの定数倍の和で表現できる。

以下の84x84の行列では、最初41行に F_{i} \times (i)^j、次の41行に F_{i-1} \times (i-1)^jの係数、83行目に S_{i}=\sum_{j=1}^i F_j \times j^K、84行目に S_{i-1}の係数を配置している。
そしてその行列を(N-1)乗して F_{1} \times (1)^j = 1、 F_{0} \times (0)^j = 0(ただしj==0の時は1)を要素に持つベクトルと掛け合わせて答えを求めた。

例えばK==4の時の行列は以下のようになる。
84x84全部は大きいので省略して、先頭6行がF[i]*i^j、次の6行がF[i-1]*(i-1)^j、最後2行がS[i]とS[i-1]を表している。

   1    0    0    0    0    0    1    0    0    0    0    0    0    0
   1    1    0    0    0    0    2    1    0    0    0    0    0    0
   1    2    1    0    0    0    4    4    1    0    0    0    0    0
   1    3    3    1    0    0    8   12    6    1    0    0    0    0
   1    4    6    4    1    0   16   32   24    8    1    0    0    0
   0    0    0    0    0    0    0    0    0    0    0    0    0    0
   1    0    0    0    0    0    0    0    0    0    0    0    0    0
   0    1    0    0    0    0    0    0    0    0    0    0    0    0
   0    0    1    0    0    0    0    0    0    0    0    0    0    0
   0    0    0    1    0    0    0    0    0    0    0    0    0    0
   0    0    0    0    1    0    0    0    0    0    0    0    0    0
   0    0    0    0    0    1    0    0    0    0    0    0    0    0
   1    4    6    4    1    0   16   32   24    8    1    0    1    0
   0    0    0    0    0    0    0    0    0    0    0    0    1    0

計算量は、行列の大きさがO(K)なので累乗処理でO(logN K^3)。

ll N,K;
ll mo=1000000007;

const int CN=1001;
ll C[CN][CN];

const int MAT=85;
ll G[MAT][MAT],G2[MAT][MAT];

void powmat(ll p,int n,ll a[MAT][MAT],ll r[MAT][MAT]) {
	int i,x,y,z;
	ll a2[MAT][MAT];
	FOR(x,n) FOR(y,n) a2[x][y]=a[x][y];
	FOR(x,n) FOR(y,n) r[x][y]=0;
	FOR(i,n) r[i][i]=1;
	while(p) {
		ll h[MAT][MAT];
		if(p%2) {
			FOR(x,n) FOR(y,n) {
				h[x][y]=0;
				FOR(z,n) h[x][y] += (r[x][z]*a2[z][y]) % mo;
				h[x][y] %= mo;
			}
			FOR(x,n) FOR(y,n) r[x][y]=h[x][y];
		}
		FOR(x,n) FOR(y,n) {
			h[x][y]=0;
			FOR(z,n) h[x][y] += (a2[x][z]*a2[z][y]) % mo;
			h[x][y] %= mo;
		}
		FOR(x,n) FOR(y,n) a2[x][y]=h[x][y];
		p>>=1;
	}
}

void solve() {
	int f,i,j,k,l,x,y;
	
	cin>>N>>K;
	if(N==1) return _P("1\n");
	
	FOR(x,CN) C[x][0]=C[x][x]=1;
	for(i=1;i<CN;i++) for(j=1;j<i;j++) C[i][j]=(C[i-1][j-1]+C[i-1][j])%mo;
	
	FOR(i,K+1) {
		FOR(j,i+1) {
			G[i][j]=C[i][j];
			FOR(l,j+1) G[i][l+41]+=C[i][j]*C[j][l],G[i][l+41]%=mo;
		}
	}
	memmove(G[82],G[K],sizeof(G[K]));
	for(i=41;i<82;i++) G[i][i-41]=1;
	G[82][82]=1;
	G[83][82]=1;
	
	powmat(N-1,84,G,G2);
	FOR(i,41) ret+=G2[82][i];
	ret+=G2[82][82]+G2[82][41];
	
	cout << ret%mo << endl;
	
}

まとめ

行列表現までは思いついたけど、この行列は思いつかなかった…。