2021牛客暑期多校训练营4 B. Sample Game

题意：每次从1到n中抽一个数，概率为p[i]，如果抽出的数比上一个小，则停止抽取，否则继续抽，问抽取次数平方的期望

我们要求的式子长这样

\[ans=\sum\limits_{i=1}^{\infty} P(len=i) i^2 \]

套路地转化为

\[\sum\limits_{i=1}^{\infty}(P(len\ge i-1)-P(len\ge i ))i^2 \]

做变量代换得到

\[\sum\limits_{i=0}^\infty P(len\ge i)(2i+1) \]

也即

\[2\sum\limits_{i=1}^\infty P(len\ge i)i+\sum\limits_{i=0}^\infty P(len\ge i) \]

构造概率生成函数

\[f(x)=\sum\limits_{i=0}^\infty P(len\ge i ) x^i \]

求导得

\[f^{'}(x)=\sum\limits_{i=1}^\infty P(len\ge i )x^{i-1}i \]

那么答案就是

\[ans=2f^{'}(1)+f(1) \]

考虑\(f(x)\)的实际含义，即每个数出现概率的生成函数相乘

\[f(x)=\prod\limits_{i=1}^n(1+P_ix+P_i^2x^2+...) \]

即为

\[f(x)=\prod\limits_{i=1}^n \dfrac{1}{1-P_ix} \]

对\(f(x)\)求导（可以采用两边取对数或直接求导）

\[f^{'}(x)=f(x)\sum\limits_{i=1}^n \dfrac{P_i}{1-P_ix} \]

至此可以O(n)求出答案，加上求逆元复杂度O(nlogn)

#include<bits/stdc++.h>
#define int long long
using namespace std;

const int MOD = 998244353;
const int MAXN = 1005;

int rd(){
  int ret=0,f=1;char c;
  while(c=getchar(),!isdigit(c))f=c=='-'?-1:1;
  while(isdigit(c))ret=ret*10+c-'0',c=getchar();
  return ret*f;
}

int qpow(int x,int y=MOD-2){
  int ret=1,base=x;
  while(y){
    if(y&1) ret=ret*base%MOD;
    base=base*base%MOD;
    y>>=1;
  }
  return ret;
}

int p[MAXN],n,sum;

signed main(){
  n=rd();
  for(int i=1;i<=n;i++)
    sum+=(p[i]=rd());
  sum=qpow(sum);
  for(int i=1;i<=n;i++)
    p[i]=p[i]*sum%MOD;
  int f=1,df=0;
  for(int i=1;i<=n;i++)
    f=f*qpow(MOD+1-p[i])%MOD;
  for(int i=1;i<=n;i++)
    df=(df+p[i]*qpow(MOD+1-p[i]))%MOD;
  df=df*f%MOD;
  cout<<(2*df+f)%MOD;
}