HackerRank HourRank 21 : B. Sam's Numbers

Cが詰め切れなかったがなんとかTシャツ圏内に入った。
https://www.hackerrank.com/contests/hourrank-21/challenges/sams-numbers

問題

パラメータS,M,Dが与えられる。
以下の条件を満たす数列はいくつあるか。

• 各要素は1～Mの範囲の整数である。
• 隣り合う要素の差の絶対値はD以下である。
• 総和はSである。

解法

f(n,k) := 数列の総和がnで末尾がkであるような物の組み合わせ
とすると、以下のようになる。

• f(k,k) := 1≦k≦Mなら1、それ以外0
• 

求めたいのはf(S,*)の総和である。
まずS≦Mの場合はDPで求めてしまおう。
それ以外の場合、f(n,*)の値はf(n-k,*) (1≦k≦M)で決まるので、行列累乗で漸化式を処理しよう。
(M*M)次正方行列の累乗をしていけばよい。

ll S,M,D;
ll mo=1000000009;

const int MAT=100;
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]=(x==y);
	while(p) {
		ll h[MAT][MAT];
		if(p%2) {
			FOR(x,n) FOR(y,n) h[x][y]=0;
			FOR(x,n) FOR(z,n) FOR(y,n) h[x][y] += (r[x][z]*a2[z][y]) % mo;
			FOR(x,n) FOR(y,n) r[x][y]=h[x][y]%mo;
		}
		FOR(x,n) FOR(y,n) h[x][y]=0;
		FOR(x,n) FOR(z,n) FOR(y,n) h[x][y] += (a2[x][z]*a2[z][y]) % mo;
		FOR(x,n) FOR(y,n) a2[x][y]=h[x][y]%mo;
		p>>=1;
	}
}

ll F[21][11];
ll H[100];

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>S>>M>>D;
	
	FOR(i,M+1) F[i][i]=1;
	for(i=1;i<=21;i++) {
		for(j=1;j<=M;j++) {
			for(x=1;x<=M;x++) if(abs(j-x)<=D) (F[i+x][x] += F[i][j])%=mo;
		}
	}
	
	ll ret=0;
	if(S<=M) {
		FOR(i,M+1) ret+=F[S][i];
		cout<<ret%mo<<endl;
		return;
	}
	
	for(i=1;i<M;i++) {
		for(j=1;j<=M;j++) G[M*i+j-1][M*(i-1)+j-1]=1;
	}
	for(i=1;i<=M;i++) {
		for(j=1;j<=M;j++) if(abs(i-j)<=D) {
			G[i-1][(i-1)*M+(j-1)]=1;
		}
	}
	
	
	powmat(S-M,M*M,G,G2);
	FOR(x,M) FOR(y,M) H[x*M+y]=F[M-x][y+1];
	FOR(x,M) FOR(y,M*M) ret+=G2[x][y]*H[y]%mo;
	cout<<ret%mo<<endl;
	
}

まとめ

(M*M)次正方行列の累乗をするので計算量はO(M^6*logS)か。