AtCoder ABC #239 (デンソークリエイトプログラミングコンテスト2022) : Ex - Dice Product 2

なんとか解けてよかったね。
https://atcoder.jp/contests/abc239/tasks/abc239_h

問題

N面のダイスがある。
今変数Xの初期値を1とする。
ダイスを振って、そのたびその目の分だけXに掛け算することを考える。
XがMを超えるまでのダイスロール回数の期待値を求めよ。

解法

f(n) := 今X=nであるとき、XがM超えるまでの追加ダイスロール回数の期待値
とするとf(1)が解となる。ただ、このまま考えると状態がO(M)個あるし、状態遷移も平均O(logM)あるので全体でO(MlogM)かかる。

そこでちょっと考え方を変える。
g(m) := 今あるXに対し、最小であとm倍すればXがMを超えるような状態における、XがM超えるまでの追加ダイスロール回数の期待値
n≦√Mの時はf(n)、n＞√Mの時はg(ceil*1を考えるようにすると、状態数がO(√M)で済む。
状態遷移も、愚直にやると各状態でO(√M)程度かかるが、遷移先が同じダイスの目はまとめて処理するようにしよう。

ll N,M;
const ll mo=1000000007;

ll F[101010];
ll G[101010];

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

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>M;
	M++;
	int m=1;
	while(m*m<M) m++;
	
	ll re=modpow(N-1);
	ll si=N*modpow(N-1)%mo;
	
	for(i=2;i<=m;i++) {
		for(j=2;j<=N;j++) {
			x=(i+j-1)/j;
			if(x==1) break;
			k=min((int)N,(i-1)/(x-1));
			
			F[i]+=F[(i+j-1)/j]*(k-j+1)%mo;
			j=k;
		}
		F[i]=(si+F[i]%mo*re)%mo;
	}
	
	for(i=m;i>=1;i--) {
		for(j=2;j<=N&&1LL*i*j<M;j++) {
			if(i*j<=m) {
				G[i]+=G[i*j];
			}
			else {
				x=(M+i*j-1)/(i*j);
				if(x==1) break;
				k=min(N,(M-1)/(x-1)/i);
				G[i]+=F[(M+i*j-1)/(i*j)]*(k+1-j)%mo;
				j=k;
			}
		}
		G[i]=(si+G[i]%mo*re)%mo;
	}
	cout<<G[1]<<endl;
}

まとめ

平方根で分けるの、ちゃんと思いつけて良かった。

*1:M+1)/n