题目大意
给定生成 $1$ 到 $n$ 数字的概率分布，不断生成随机数 $x$ 直到 $x$ 不是已经生成过的最大的数停止，求生成次数平方的期望。
分析
设期望的随机次数为 $f_x = E(x)$，我们需要计算的次数为 $g_x = E(x^2)$。
关于 $f_i$ 的 DP 是显然的，计算有
$$ f_x = \sum_{i = 1}^{x - 1} p_i + \sum_{i = x}^n p_i(1 + f_i) = 1 + \sum_{i = x}^n p_i f_i $$
容易观察到
$$ f_{x + 1} =(1 - p_x) f_x \Rightarrow f_x = \prod_{i = x}^n \frac{1}{1 - p_x} $$
接下来需要一个套路
$$ E((x + 1)^2) = E(x^2) + 2 E(x) + 1 $$
类似的可以推得
$$ g_x = \sum_{i = 1}^{x - 1} p_i + \sum_{i = x}^n p_i(g_i + 2 f_i + 1) = 1
- \sum_{i = x}^n p_i g_i + 2 \sum_{i = x}^n p_i f_i = \sum_{i = x}^n p_i g_i + 2 f_x - 1 $$
最终答案即是
$$ ans = \sum_{i = 1}^n p_i(g_i + 2 f_i + 1) = g_1 $$
至此，倒着递推已经可以线性求解了。但是我们还可以继续优化，逐项相减有
$$ g_{x + 1} - g_x = 2(f_{x + 1} - f_x) - p_x g_x = - 2 p_x f_x - p_x g_x $$
即
$$ \frac{g_x}{f_x} - \frac{g_{x + 1}}{f_{x + 1}} = \frac{2}{1 - p_x} - 2 $$
因此
$$ {\rm ans} = g_1 = f_1 \left( \frac{g_n}{f_n} - 2(n - 1) + 2 \sum_{i = 1}^{n - 1} \frac{1}{1 - p_x} \right) = \left( \prod_{i = x}^n \frac{1}{1 - p_x} \right) \left( 1 + 2 \sum_{i = 1}^n \frac{1}{1 - p_x} - 2 n \right) $$
至此，我们可以 $O(n)$ 的解决问题。
#define ACM_MOD 998244353
const ll mod = ACM_MOD;
#include "template/basic/qpow.hpp"
#include "template/basic/inv.hpp"
ll p[10086];
int main() {
ll n = rr(), sum = 0;
for (int i = 1; i <= n; i++)
p[i] = rr(), sum = (sum + p[i]) % mod;
sum = inv(sum);
for (int i = 1; i <= n; i++) {
p[i] = p[i] * sum % mod;
p[i] = inv(mod + 1 - p[i]);
}
ll ans = 1, w = mod - n;
for (int i = 1; i <= n; i++) {
ans = ans * p[i] % mod;
w = (w + p[i]) % mod;
}
ans = ans * (1 + 2 * w) % mod;
printf("%lld\n", ans);
return 0;
}
我们经过了很长的化简过程才导出这一结果，实际上生成函数更为快捷。
设 $f(x)$ 是生成长度为 $i$ 的非递减序列的生成函数，即 $P(len > i)$，可以推出
$$ f(x) = \prod_{i = 1}^n \frac{1}{1 - p_i x} $$
而我们需要求
$$ \sum_{i = 1}^{\infty} (f_{i - 1} - f_i) i^2 = \sum_{i = 0}^{\infty} f_i (2 i + 1) = 2 f’(1) + f(1) $$
化简即可得到上式。