Codeforces ECR #078 : F. Cards

これのお陰で後日AtCoderのコンテストで解法の目途が立ったよね。
https://codeforces.com/contest/1278/problem/F

問題

M枚のカードがあり、1枚がジョーカーである。
これらのカードから1枚をランダムに取り、戻すことをN回繰り返す。

ジョーカーが出た回数をXとすると、この時のスコアを定数Kを用いてX^Kとする。
スコアの期待値を求めよ。

解法

N,Mが大きいので、O(N)やO(M)かかる解法は使えない。
Kがさほど大きくないことに着目しよう。
ある数aが良いとは、a回目の施行でジョーカーが出たことを示す。

もしK=2の時、2要素のタプル(a1,a2)を考える。
全要素が良い値であるようなタプルの総数がスコアの値となる。
Kが3以上の場合も同様。

K要素中m個uniqueであるような値の選び方をdp(K,m)とする。
m個がいずれも良い値ならよいので、dp(K,m)*(1/M)^mがスコアに計上される。

int N,M,K;
const ll mo=998244353;


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

ll dp[5050][5050];

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>M>>K;
	
	dp[0][0]=1;
	FOR(i,K) {
		FOR(j,i+1) {
			// same
			dp[i+1][j]+=dp[i][j]*j%mo;
			dp[i+1][j+1]+=dp[i][j]*(N-j)%mo;
		}
	}
	ll ret=0;
	ll revm=modpow(M);
	ll mi=1;
	for(i=1;i<=K;i++) {
		mi=mi*revm%mo;
		ret+=dp[K][i]*mi%mo;
	}
	
	cout<<ret%mo<<endl;
}

まとめ

ECRの最終問にしてはコード量が短い。