yukicoder : No.1516 simple 門松列 problem Re:MASTER

典型ではあるけどちょっとめんどい。
https://yukicoder.me/problems/no/1516

問題

正整数N,Kが与えられる。
0～(K-1)の値を取るN要素の数列にはK^N通りが考えられる。
このうち門松列列をなすものの数と、値の総和を求めよ。

解法

f(n,a,b) := n要素目まで値を決めたとき、末尾2要素がa,bである門松列列であるような組み合わせ
g(n,a,b) := 上記門松列列の値の総和

とすると、(a,b,c)が門松列をなす場合、
f(n+1,b,c) += f(n,a,b)
g(n+1,b,c) += g(n,a,b) + c*f(n,a,b)
と遷移する。あとは2*K*K次の行列乗算に持ち込もう。

int N,K;
const ll mo=998244353;
const int MAT=163;
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;
}


void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>K;
	Mat A;
	FOR(x,K) FOR(y,K) FOR(i,K) {
		if(x!=i&&((x<y&&i<y)||(x>y&&i>y))) {
			A.v[y*K+i][x*K+y]=1;
			A.v[K*K+y*K+i][K*K+x*K+y]=1;
			A.v[K*K+y*K+i][x*K+y]=i;
		}
	}
	A=powmat(N-2,A,2*K*K+1);
	ll ret1=0,ret2=0;
	FOR(x,K) FOR(y,K) if(x!=y) {
		FOR(i,K) FOR(j,K) {
			ret1+=A.v[i*K+j][x*K+y];
			ret2+=A.v[K*K+i*K+j][x*K+y];
			ret2+=A.v[K*K+i*K+j][K*K+x*K+y]*(x+y);
		}
	}
	cout<<ret1%mo<<" "<<ret2%mo<<endl;
	
}

まとめ

ギリギリ間に合ってよかった。