AtCoder ABC #256 (東京海上日動プログラミングコンテスト2022) : G - Black and White Stones

0x100回目おめでとうございます。
https://atcoder.jp/contests/abc256/tasks/abc256_g

問題

正整数N,Dが与えられる。
1辺の長さがDの正N角形を考える。
頂点から初めて、距離1ごとに白か黒の石を置くと、各辺上に(D+1)個の石が置かれる。

各辺上に置かれた白石の数が一致する石の色の選び方は何通りか。

解法

Dが小さいので、白石の数を総当たりしよう。
問題は多角形の頂点上の石で、この石の状態は2辺に寄与するためややこしい。
そこで、各頂点について、石が黒の場合と白の場合をそれぞれ
f(n,a,b) := ある頂点から初めて、n個目までの頂点上の石の色を決めたとき、aを0個目の頂点上の石の白黒、bをn個目の頂点上の石の白黒

とすると、f(n,a,b)→f(n+1,a,c)の遷移は、前処理で二項係数をO(1)で求めらえるようにしておけば、それぞれO(1)で求められる。
あとは行列累乗を行い、f(N,白,白)とf(N,黒,黒)の和を答えよう。

ll N,D;
const ll mo=998244353;

const int MAT=2;
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;
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>D;
	ll ret=1;
	for(int d=1;d<=D+1;d++) {
		Mat A;
		A.v[0][0]=comb(D-1,d);
		A.v[0][1]=comb(D-1,d-1);
		A.v[1][0]=comb(D-1,d-1);
		A.v[1][1]=comb(D-1,d-2);
		A=powmat(N,A);
		(ret+=A.v[0][0]+A.v[1][1])%=mo;
	}
	cout<<ret<<endl;
}

まとめ

解法はすぐ思いつけた。