题目链接
https://atcoder.jp/contests/abc352/tasks/abc352_g
题目大意
总共有 \(N\) 种不同颜色的袜子，第 \(i\) 种颜色的袜子有 \(A_i\) 只。现在高桥要按照以下规则决定早上穿哪一双袜子：
从当前剩余的袜子中等概率随机且不放回地抽取一只袜子。
若抽出的袜子颜色与之前某次抽出的袜子颜色相同，则停止操作并穿这种颜色的袜子；
若抽出的袜子颜色与之前每一次抽出的颜色均不同，则进行下一次抽取操作。
求高桥抽取袜子次数的期望值，对 \(998244353\) 取模。
数据范围
- \(1 \le N \le 3\times 10^{5}\)
- \(2 \le A_i \le 3000\)
解题思路
设 \(f_k\) 代表抽了 \(k\) 次袜子之后还没有凑出来一对相同颜色袜子的概率，那 \(f_k - f_{k + 1}\) 就是第 \(k + 1\) 次抽袜子凑出来了一对相同颜色袜子的概率，设 \(tot = \sum A_i\)，那么期望次数就是 \(\displaystyle{\sum_{k = 0}^{n}} (f_k - f_{k + 1}) \times (k + 1)\)。接下来考虑怎么求出 \(f\)。
因为每种袜子抽到不能超过 \(1\) 次，考虑枚举抽到了 \(p_1, p_2, \cdots, p_k\) 这些种类的袜子，那么有 $$ \[\begin{aligned} f_k &= \sum\limits_{\{p_1, p_2, \cdots, p_k\} \subset \{1, 2, \cdots, n\}} k! \cdot \frac{A_{p_1}}{tot} \cdot \frac{A_{p_2}}{tot - 1} \cdots \frac{A_{p_k}}{tot - (k - 1)} \\ &= \frac{k!(tot - k)!}{tot!} \sum\limits_{\{p_1, p_2, \cdots, p_k\} \subset \{1, 2, \cdots, n\}} A_{p_1} A_{p_2}\cdots A_{p_k} \end{aligned}\]$$ 令 \(g_k = \displaystyle\sum\limits_{\{p_1, p_2, \cdots, p_k\} \subset \{1, 2, \cdots, n\}} A_{p_1} A_{p_2}\cdots A_{p_k}\)，注意到 \(g_k\) 的含义就是从 \(n\) 个里面选 \(k\) 个所有方案的乘积之和。
可以把序列分为左右两部分，两部分的 \(g\) 分别求出来，做一遍卷积就得到合起来的 \(g\) 了，使用分治即可，时间复杂度 \(O(n\log^{2}n)\)。