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

题目大意

玩一场游戏，每一轮在 $[1,n]$ 中随机生成一个整数 $x$ ，生成 $x$ 的概率为 $p_x$ 。若 $x$ 不小于之前生成的任何数，则游戏继续并进入下一轮，否则游戏结束。假如说最后游戏共进行了 $len$ 轮，则游戏的得分为 $len^2$ 。求游戏分数的期望值 $E(len^2)$ ，答案 $\pmod{998244353}$ 输出。

输入为 $n$ 个数 $w_1,w_2,\cdots,w_n$ ， $p_i=\frac{w_i}{\sum\limits_{j=1}^n w_j}$ 。

思路

游戏生成的数字序列一定是类似 $1,1,\cdots,1,2,2\cdots,3,3\cdots$ ，即 $[1,n]$ 的数字由小到大依次产生，每个数字出现的次数为 $[0,+\infty)$ 。

构造母函数 $f(x)=\prod\limits_{i=1}^n\sum\limits_{m=0}^\infty p_i^mx^m\ (1)$ ，对于每一项 $\sum\limits_{m=0}^\infty p_i^mx^m$ ，指数 $m$ 表示数字 $i$ 出现的次数， $p_i^m$ 就是 $i$ 连续出现 $m$ 次的概率。每一个数的出现次数都对应多项式 $\sum\limits_{m=0}^\infty p_i^mx^m$ ，将这 $n$ 个多项式进行卷积运算，即为 $f(x)$ 。卷积的结果为 $f(x)=\sum\limits_{m=0}^\infty a_mx^m$ ， $m$ 为生成序列的长度，系数 $a_m$ 自然是生成长度为 $m$ 的非递减序列的概率。仔细思考， $a_m$ 同时也是游戏进行轮次超过 $m$ 次的概率，即 $f(x)=\sum\limits_{m=0}^\infty P(len>m)x^m\ (2)$ ；因为游戏已经进行了 $m$ 轮，且可以继续进行，游戏要进行超过 $m$ 轮等价于生成长度为 $m$ 的非递减序列。

根据数学期望的定义，枚举最后游戏的轮次 $m$ 。

$\begin{aligned} E(len^2)&=\sum\limits_{m=1}^\infty [P(len>m-1)-P(len>m)]m^2\\&=P(len>0)+\sum\limits_{m=1}^\infty P(len>m)[(m+1)^2-m^2]\\&=1+\sum\limits_{m=1}^\infty P(len>m)(2i+1)\\&=\sum\limits_{m=0}^\infty P(len>m)(2i+1) \end{aligned}$

根据 $(2)$ 可以发现， $E(len^2)=2f'(1)+f(1)$ 。我们需要用 $(1)$ 来求 $f'(1)$ 和 $f(1)$ 。首先将 $f(x)$ 化为其封闭形式，即 $f(x)=\prod\limits_{i=1}^n\frac{1}{1-p_ix}\ (3)$ 。将 $x=1$ 带入 $(3)$ ，容易得到 $f(1)=\prod\limits_{i=1}^n \frac{1}{1-p_i}$ 。将 $(3)$ 两边取自然对数， $\ln f(x)=\sum\limits_{i=1}^n\ln\frac{1}{1-p_ix}$ ；两边对 $x$ 求导， $\frac{f'(x)}{f(x)}=\sum\limits_{i=1}^n\frac{p_i}{1-p_ix}$ ；将 $x=1$ 带入，得 $f'(1)=f(1)\sum\limits_{i=1}^n\frac{p_i}{1-p_i}$ 。将 $p_i=\frac{w_i}{\sum\limits_{j=1}^n w_j}$ 代入，并记 $sum=\sum\limits_{j=1}^n w_j$ ，有 $f(1)=\prod\limits_{i=1}^n\frac{sum}{sum-w_i}$ ， $f'(1)=f(1)\sum\limits_{i=1}^n\frac{w_i}{sum-w_i}$ 。

求 $f(1)$ 和 $f'(1)$ ，时间复杂度 $O(n)$ ，最后 $O(1)$ 代入求 $E(len^2)$ 。

代码

#include<iostream>
using namespace std;
typedef long long ll;
const int MOD = 998244353;
const int MAX_SIZE = 105;
int n, w[MAX_SIZE];
ll QuickPower(ll a, ll b)
{
    ll res = 1;
    while (b)
    {
        if (b & 1)
        {
            res = res * a % MOD;
        }
        a = a * a % MOD;
        b >>= 1;
    }
    return res;
}
int main()
{
    cin >> n;
    int sum = 0;
    for (int i = 1;i <= n;i++)
    {
        cin >> w[i];
        sum += w[i];
    }
    // 求 f(1)
    int a = 1;
    for (int i = 1;i <= n;i++)
    {
        a = 1ll * a * sum % MOD * QuickPower(sum - w[i], MOD - 2) % MOD;
    }
    // 求 f'(1)
    int b = 0;
    for (int i = 1;i <= n;i++)
    {
        b = (b + 1ll * w[i] * QuickPower(sum - w[i], MOD - 2) % MOD) % MOD;
    }
    b = 1ll * b * a % MOD;
    // 输出答案
    cout << (2ll * b % MOD + a) % MOD << endl;
    return 0;
}