[生成函数]2021牛客第四场 B Sample Game

题目

有一个随机数生成器，每次生成$i$的概率为$p_i$,如果现在生成的数不小于之前生成的所有数，继续随机，反之，记上分数为$n^2$,$n$为已经生成的数的个数

思路

设最长答案长度是$len$,
设生成函数，设它的系数是$len>i$的概率$P(len>i)$，得:$f(x)={\sum }_{i=0}^{\infty }P(len>i)x^i$

对于一种情况，设第$i$个不同的数出现过$cn{t}_i$次，那么该种情况的概率为:${\prod }_{i=1}^n{p}_i^{cn{t}_i}$

所以我们可以枚举每一个数可能出现的次数，它的生成函数为：
$g_i(x)={\sum }_{j=0}^{\infty }{p}_i^j{x}^j=\frac{1}{1-p_{i}x}$

把所有可能的次数乘起来，每一个结果都对应不同的情况：
$f(x)={\prod }_{i=1}^n{g}_i(x)$
带入上方公式:$f(x)={\prod }_{i=1}^n\frac{1}{1-p_{i}x}$

长度等于$i$时的概率$P(len=i)=P(len>i-1)-P(len>i)$
答案$E（len=i）=P(len=i)\ast i^2$
所有答案为：
${\sum }_{i=1}^{\infty }(P(len>i-1)-P(len>i))\ast i^2$
$=...+P(len>i-1)i^2-P(len>i)i^2+P(len>i)\ast (i+1{)}^2-P(len>i+1)(i+1{)}^2+...$
$={\sum }_{i=0}^{\infty }(P(len>i))\ast ((i+1{)}^2-i^2)$
$（看到有地方再说P（len>0）的时候的问题，i=0的时候也满足这个式子）$
$={\sum }_{i=0}^{\infty }(P(len>i))\ast (2i+1)$

$f(x)={\sum }_{i=1}^{\infty }P(len>i)x^i$
$f^{\prime }(x)=i\ast {\sum }_{i=1}^{\infty }P(len>i)x^{i-1}$
$f(1)={\sum }_{i=1}^{\infty }P(len>i)$
$f^{\prime }(1)=i\ast {\sum }_{i=1}^{\infty }P(len>i)$

${\sum }_{i=1}^{\infty }(P(len>i-1)-P(len>i))\ast i^2$
$=...+P(len>i-1)i^2-P(len>i)i^2+P(len>i)\ast (i+1{)}^2-P(len>i+1)(i+1{)}^2+...$

代入原式子：${\sum }_{i=1}^{\infty }(P(len>i))\ast (2i+1)=2\ast f^{\prime }(1)+f(1)$

$f(1)={\prod }_{i=1}^n\frac{1}{1-p_i}$
$f^{\prime }(x)=f(x){\sum }_{i=1}^n\frac{p_i}{1-p_{i}x}$
$f^{\prime }(1)=f(1){\sum }_{i=1}^n\frac{p_i}{1-p_i}$

代码

#include <bits/stdc++.h>
#define int long long
using namespace std;
typedef long long ll;
int w[1000],p[1000];
const int mod =998244353;
int quick_pow(int a,int b){
	int res=1;
	while(b){
		if(b&1) res=res*a%mod;
		b>>=1;
		a=a*a%mod;
	}
	return res;
}
signed main(){
    ios::sync_with_stdio(0); cin.tie(0); cout.tie(0);
    int n;
    cin>>n;
    int ans=0;
    for(int i=1;i<=n;i++){
    	cin>>w[i];
    	ans+=w[i];
	}
	int inv=quick_pow(ans,mod-2)%mod;
	int pp=0,f1=1;
	for(int i=1;i<=n;i++){
		p[i]=w[i]*inv%mod;
		f1=(f1*quick_pow((1-p[i]+mod),mod-2))%mod;
		
		pp=(pp+p[i]*quick_pow(1-p[i]+mod,mod-2)%mod)%mod;
	}
	

	cout<<(f1+2*f1*pp%mod)%mod<<endl;
    return 0;
}