kmjp's blog

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

yukicoder : No.2883 K-powered Sum of Fibonacci

なるほど。
https://yukicoder.me/problems/no/2883

問題

整数N,Kが与えられる。
フィボナッチ数列F(n)について、F(1)^K+F(2)^K+....+F(N)^Kを求めよ。

F(n+1)^K = (F(n)+F(n-1))^K より、F(n+1)^KはF(n)^m*F(n-1)^(K-m)の線形和で求められる。
また、F(n+1)^m*F(n)^(K-m)についても、F(n)^l*F(n-1)^(K-l)の線形和で計算できる。
あとは行列累乗に持ち込み、総和の項を合わせて(K+2)次の正方行列のN乗を求めればよい。

const ll mo=998244353;
const int MAT=102;
struct Mat { ll v[MAT][MAT]; Mat(){ZERO(v);};};

Mat mulmat(Mat& a,Mat& b,int n=MAT) {
	ll mo2=4*mo*mo;
	int x,y,z; Mat r;
	FOR(x,n) FOR(y,n) r.v[x][y]=0;
	FOR(x,n) FOR(z,n) FOR(y,n) {
		r.v[x][y] += a.v[x][z]*b.v[z][y];
		if(r.v[x][y]>mo2) r.v[x][y] -= mo2;
	}
	FOR(x,n) FOR(y,n) r.v[x][y]%=mo;
	return r;
}

Mat powmat(ll p,Mat a,int n=MAT) {
	int i,x,y; Mat r;
	FOR(x,n) FOR(y,n) r.v[x][y]=0;
	FOR(i,n) r.v[i][i]=1;
	while(p) {
		if(p%2) r=mulmat(r,a,n);
		a=mulmat(a,a,n);
		p>>=1;
	}
	return r;
}

ll comb(ll N_, ll C_) {
	const int NUM_=400001;
	static ll fact[NUM_+1],factr[NUM_+1],inv[NUM_+1];
	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 N,K;

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>K;
	Mat A;
	FOR(y,K+1) FOR(x,K+1) if(x<=K-y) {
		A.v[y][x]=comb(K-y,K-x-y);
		if(y==0) A.v[K+1][x]=comb(K-y,K-x-y);
	}
	A.v[K+1][K+1]=1;
	A=powmat(N-1,A,K+2);
	
	cout<<A.v[K+1][0]+1<<endl;

	
}

まとめ

F(n+1)^m*F(n)^(K-m)を順次計算していくってのが賢いな。