Codeforces #641 Div1 D. Slime and Biscuits

コードは短いんだよな。
https://codeforces.com/contest/1349/problem/D

問題

N人の人がいて、初期状態でそれぞれビスケットA[i]個を持っている。
この状態で以下を繰り返す。

• 全ビスケットのうち、1つランダムで選択される。
• そのビスケットの所有者は、残り(N-1)人のいずれかに等確率でランダムに移動する。

全ビスケットが1人に集まるまでの処理回数の期待値を求めよ。

解法

F(x) := 今x個ビスケットを持ってる人が、最後全ビスケットを得るまでにかかる時間の期待値
E(i) := i番目の人が全ビスケットを得る確率と、かかる時間の期待値の積
とする。
求めたいのはE(i)の総和で、E(i)=F(A[i])となる。
詳細はEditorialを参照してもらうとして、N人に関する期待値の式の総和を足すと、解は

となる。
G(x) := 今x個ビスケットを持ってる人が、x+1個のビスケットを得るまでにかかる時間の期待値
とすると、G(x)を小さい順に求めたうえで、F(x)=G(x)+F(x+1)を大きい順に求めていくとF(x)が求まる。

int N,M;
int A[303030];
const ll mo=998244353;

ll F[303030];


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;
	FOR(i,N) {
		cin>>A[i];
		M+=A[i];
	}
	F[0]=N-1;
	for(x=1;x<M;x++) {
		ll a=(M-x)*modpow(M)%mo*modpow(N-1)%mo;
		ll b=1+x*F[x-1]%mo*modpow(M)%mo;
		F[x]=b*modpow(a)%mo;
	}
	
	for(i=M-1;i>=0;i--) F[i]=(F[i]+F[i+1])%mo;
	
	ll ret=0;
	FOR(i,N) ret+=F[A[i]];
	ret=(ret%mo-F[0]*(N-1)%mo+mo)%mo;
	cout<<ret*modpow(N)%mo<<endl;
}

まとめ

コードはともかく式変形がかなり難しい。