概率dp AtCoder Beginner Contest 243 F

AtCoder Beginner Contest 243 F

题意

n种彩票，每种彩票i中奖概率为$p_i$ ，求抽k次后恰好m种彩票中奖的概率（简化的概率，题目中的概率还要用个逆元）

思路

超几何分布，设抽k次后每种彩票i中奖个数为$c_i$ ，则概率为
$P=C_k^{c_1}*p_1^{c_1}*C_{k-c_1}^{c_2}*p_2^{c_2}*C_{k-c_1-c_2}^{c_3}*p_3^{c_3}...*C_{k-\sum_{i=1}^{n-1}}^{c_n}*p_n^{c_n}$,化简之后就变成$P=k!*\frac{\prod_{i=1}^n p_i^{c_i}}{\prod_{i=1}^n c_i!}$
这样就清晰很多了，设状态$f[i][j][k]$为当前第i种彩票，已经由j种彩票被抽中，抽了k次的概率；
那么有 $f[i][j][k]=\sum f[i-1][j-(c!=0)][k-c]*\frac{p_i^c}{c!}$，即当前第i种彩票买了c张，从之前的状态转移过来

代码

const int N=55,inf=0x3f3f3f3f,mod= 998244353;
int n,m,K;
int f[N][N][N],fac[N],p[N],inv_f[N];

int q_pow(int a,int b)
{
	int res=1;
	while(b)
	{
		if(b&1) res=res*a%mod;
		a=a*a%mod;
		b/=2;
	}
	return res%mod;
}

void init(int n)
{
	fac[0]=inv_f[0]=1;
	for(int i=1;i<=n;i++) fac[i]=fac[i-1]*i%mod;
	inv_f[n]=q_pow(fac[n],mod-2)%mod;
	for(int i=n-1;i>=1;i--) 
	{
		inv_f[i]=inv_f[i+1]*(i+1)%mod;
	}
}

void solve() 
{
	cin>>n>>m>>K;
	init(K+1);
	int sum=0;
	for(int i=1;i<=n;i++) cin>>p[i],sum+=p[i];//分母
	int inv=q_pow(sum,mod-2)%mod;
	f[0][0][0]=1;
	for(int i=1;i<=n;i++) p[i]=p[i]*inv%mod;
	for(int i=0;i<=n;i++) 
	{
	   for(int j=0;j<=i;j++) 
	   {
	      for(int k=0;k<=K;k++) 
	      {
		 for(int c=0;c+k<=K;c++) 
		 {
		    f[i+1][j+(c!=0)][k+c]=(f[i+1][j+(c!=0)][k+c]+f[i][j][k]*q_pow(p[i+1],c)%mod*inv_f[c]%mod)%mod;
		 }
	      }
	   }
	}
	int ans=f[n][m][K]*fac[K]%mod;
	cout<<ans<<endl;
}

长夜路漫漫