AtCoder ABC #299 (東京海上日動プログラミングコンテスト2023) : Ex - Dice Sum Infinity

本番は筋は合ってたな…。
https://atcoder.jp/contests/abc299/tasks/abc299_h

問題

変数Cの初期値は0とする。
6面サイコロを振り、目の分だけCを加算することを繰り返す。
正整数Rが与えられるので、C-Rが10^9の倍数になるまでのサイコロを振る回数の期待値を求めよ。

解法

以下を考える。
e(X) := 非負整数Xから、サイコロの目の数だけ値を減らすことを繰り返した場合、0以下になるまでのサイコロを振る回数の期待値と、Xが0～(-5)のそれぞれで終わる確率
これは、Xがある値に遷移する確率と、確率×サイコロを振る回数の期待値を行列累乗を使って求めることができる。

E(X) := 現在の値がXであるとき、X%10^9が0となるまでのサイコロを振る回数の期待値
とする。e(X)からいったんXが0以下になる状態を考えると、Xを10^9で割った余りを考えると、到達するXは0,10^9-1,10^9-2,10^9-4,10^9-5のいずれかである。
E(0)、E(10^9-1)、E(10^9-2)、E(10^9-3)、E(10^9-4)、E(10^9-5)について、上記e(X)を用いると6つの6元方程式ができるので、掃き出し法でそれぞれ求めればよい。

int R;
const ll mo=998244353;

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

ll modpow(ll a, ll n = mo-2) {
	ll r=1;a%=mo;
	while(n) r=r*((n%2)?a:1)%mo,a=a*a%mo,n>>=1;
	return r;
}

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;
}

pair<ll,ll> E2(int R) {
	//ちょうどR進む確率と回数
	if(R<0) return {0,0};
	
	Mat M;
	M.v[0][0]=M.v[0][1]=M.v[0][2]=M.v[0][3]=M.v[0][4]=M.v[0][5]=modpow(6);
	M.v[6][6]=M.v[6][7]=M.v[6][8]=M.v[6][9]=M.v[6][10]=M.v[6][11]=modpow(6);
	M.v[6][0]=M.v[6][1]=M.v[6][2]=M.v[6][3]=M.v[6][4]=M.v[6][5]=modpow(6);
	int i;
	FOR(i,5) M.v[i+1][i]=M.v[i+7][i+6]=1;
	M=powmat(R,M);
	if(M.v[6][0]) (M.v[6][0]*=modpow(M.v[0][0]))%=mo;
	return {M.v[0][0],M.v[6][0]};
}
vector<ll> E(int R) {
	vector<ll> ret(7);
	
	if(R<=0) {
		ret[1-R]=1;
		return ret;
	}
	
	int i,j;
	for(i=1;i<=6;i++) {
		auto p=E2(R-i);
		for(j=1;j<=6;j++) if(i-j<=0) {
			(ret[0]+=p.first*modpow(6)%mo*(p.second+1))%=mo;
			(ret[1-(i-j)]+=p.first*modpow(6))%=mo;
		}
	}
	return ret;
}


ll ma[MAT][MAT],V[MAT],RR[MAT];
ll rev(ll a, ll n = mo-2) {
	ll r=1;
	while(n) r=r*((n%2)?a:1)%mo,a=a*a%mo,n>>=1;
	return r;
}

int Gauss(int n,ll mat_[MAT][MAT],ll v_[MAT],ll r[MAT]) {
	int i,j,k;
	ll mat[MAT][MAT],v[MAT];
	memmove(mat,mat_,sizeof(mat));
	memmove(v,v_,sizeof(v));
	
	FOR(i,n) {
		if(mat[i][i]==0) {
			for(j=i+1;j<n;j++) if(mat[j][i]) break;
			if(j>=n) return i;
			FOR(k,n) swap(mat[i][k],mat[j][k]);
			swap(v[i],v[j]);
		}
		(v[i]*=rev(mat[i][i]))%=mo;
		for(k=n-1;k>=i;k--) (mat[i][k]*=rev(mat[i][i]))%=mo;
		for(j=i+1;j<n;j++) {
			v[j]=((v[j]-v[i]*mat[j][i]%mo)+mo)%mo;
			for(k=n-1;k>=i;k--) mat[j][k]=((mat[j][k]-mat[i][k]*mat[j][i]%mo)+mo)%mo;
		}
	}
	
	for(i=n-1;i>=0;i--) {
		for(j=n-1;j>i;j--) v[i]=((v[i]-mat[i][j]*v[j]%mo)+mo)%mo;
		r[i]=v[i];
	}
	return n;
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>R;
	auto ret=E(R);
	
	
	ma[0][0]=1;
	for(i=1;i<=5;i++) {
		auto p=E(1000000000-i);
		FOR(x,6) {
			ma[i][x]=p[x+1];
		}
		(ma[i][i]+=mo-1)%=mo;
		V[i]=mo-p[0];
	}
	
	Gauss(6,ma,V,RR);
	
	ll tot=ret[0];
	FOR(i,6) (tot+=ret[i+1]*RR[i])%=mo;
	cout<<tot<<endl;
}

まとめ

色々混乱してしまってもったいない。
まぁ前半グダグダだったんだけどね。