[集训队作业2018]喂鸽子

这种题的一个常用转化是，只考虑所有有效的操作（比如在这题中喂一个没饱的鸽子），假设还剩\(k\)个元素可以进行有效操作，那么进行下一次有效操作的期望次数是\(\frac{n}{k}\)。

考虑这题，我们把每一次喂到的没饱的鸽子写成一个序列，那么这个序列的长度是\(nk\)。我们先假设鸽子是按从\(1\)到\(n\)的顺序依次饱的，最后再进行一些操作计算答案。

设\(f_{i,j}\)表示考虑序列中前\(i\)个元素，\(1\)到\(j\)的鸽子饱了的概率，\(g_{i,j}\)表示考虑序列中前\(i\)个元素，\(1\)到\(j\)的鸽子饱了的概率×期望操作次数。我们考虑这样DP：当我们希望放入一个数\(x\)时，如果\(x\)没饱，那么我们先填个\(0\)上去，如果\(x\)饱了，那么我们就在前面的\(0\)中再取出\(k-1\)个全部改为\(x\)。转移时，枚举下一位填\(0\)还是填\(j+1\)。如果是填\(0\)，那么\(f'=f\times\frac{1}{n-j}\)（因为我们假设我们强制了这一位填某个数\(x\)，最终到底填什么由后面的组合数来决定），\(g'=(g+\frac{n}{n-j}\times f_{i,j})\times\frac{1}{n-j}\)。如果填\(j+1\)，那么在上述转移的基础上乘一个组合数\(\binom{i-jk}{k-1}\)。

最后答案是\(g_{nk,n}\times n!\)，乘\(n!\)是因为\(g\)表示的是概率×期望操作次数，我们强制了鸽子是按从\(1\)到\(n\)的顺序依次饱的，这个条件贡献了\(\frac{1}{n!}\)的概率。

#include<bits/stdc++.h>
using namespace std;
const int N=1e5+10;
const int mod=998244353;

int f[N][55],g[N][55],fac[N],inv[N],ifac[N],s,n,k;

int C(int n,int m) {
	if(m<0||n<m) return 0;
	return 1ll*fac[n]*ifac[m]%mod*ifac[n-m]%mod;
}

int add(int a,int b) { return a+b>=mod?a+b-mod:a+b; }
void inc(int &a,int b) { a=a+b>=mod?a+b-mod:a+b; }

int main() {
#ifndef ONLINE_JUDGE
	freopen("a.in","r",stdin);
	freopen("a.out","w",stdout);
#endif
	cin>>n>>k,s=n*k;
	fac[0]=1;for(int i=1;i<=s;i++) fac[i]=1ll*fac[i-1]*i%mod;
	inv[1]=1;for(int i=2;i<=s;i++) inv[i]=1ll*(mod-mod/i)*inv[mod%i]%mod;
	ifac[0]=1;for(int i=1;i<=s;i++) ifac[i]=1ll*ifac[i-1]*inv[i]%mod;
	f[0][0]=1;
	for(int i=0,x,y;i<s;i++)
		for(int j=0;j<n;j++)
			if(f[i][j]||g[i][j]) {
				inc(f[i+1][j],x=1ll*f[i][j]*inv[n-j]%mod);
				inc(g[i+1][j],y=1ll*add(g[i][j],1ll*n*inv[n-j]%mod*f[i][j]%mod)*inv[n-j]%mod);
				inc(f[i+1][j+1],1ll*x*C(i-j*k,k-1)%mod);
				inc(g[i+1][j+1],1ll*y*C(i-j*k,k-1)%mod);
			}
	printf("%lld\n",1ll*g[s][n]*fac[n]%mod);
	return 0;
}