题解

思路

设一种颜色出现了 $f$ 次。选择 $K$ 个糖果时，这种颜色没有被选中的概率为

$$\frac{\binom{N-f}{K}}{\binom NK}.$$

利用期望的线性性，把每种颜色至少出现一次的概率相加即可。若有 $D$ 种颜色，答案为

$$D-\frac{\sum_f h_f\binom{N-f}{K}}{\binom NK},$$

其中 $h_f$ 表示恰好出现 $f$ 次的颜色数量。

做法

平方做法可以对每个 $K$ 枚举所有非零 $h_f$。满分做法把分子同时计算。

令 $j=N-f$，定义 $A_j=h_{N-j}\,j!$，$B_t=1/t!$。则

$$\sum_f h_f\binom{N-f}{K}=\frac1{K!}\sum_{j\ge K}A_jB_{j-K}.$$

把 $A$ 反转后与 $B$ 做一次卷积，右侧求和就是卷积中下标 $N-K$ 的系数。使用模数 $998244353$ 下的 NTT 即可在 $O(N\log N)$ 内得到全部答案。

若所有颜色两两不同，任意选择 $K$ 个糖果都会得到恰好 $K$ 种颜色。若 $N\le20$，还可枚举全部子集，按子集大小累计不同颜色数。

证明

对每种颜色定义指示变量，表示所选集合是否包含该颜色。不同颜色数是所有指示变量之和，因此其期望等于各颜色出现概率之和。出现频次为 $f$ 的颜色完全未被选中的方案有 $\binom{N-f}{K}$ 种，公式成立。

阶乘展开后，固定 $K$ 的所有项恰为反转数组与逆阶乘数组卷积的第 $N-K$ 项；NTT 只改变计算方式，不改变该和。再除以 $\binom NK$ 即得到原期望，所以算法正确。

复杂度

时间复杂度为 $O(N\log N)$，空间复杂度为 $O(N)$。