Luogu P4491 [HAOI2018]染色

简单的二项式反演 + 卷积题目。

题意：对长度为 $n$ 的序列染上 $m$ 个颜色中的一个，假设恰好有 $k$ 个出现次数为 $S$，那么贡献为 $a_k$。求所有染色方案的贡献和。$n\leq 10 ^ 7$，$m\leq 10 ^ 5$。

看到恰好，二项式反演。

设 $f(k)$ 表示恰好有 $k$ 个出现次数为 $S$ 的方案数，$g(k)$ 为钦定有 $k$ 个出现次数为 $S$ 的方案数，则可以得到：

$$\begin{aligned} f(k) =& \sum_{i = k} ^ m (-1) ^ {i - k} \binom ik g(i)\\ ans =& \sum_{i = 0} ^ m f(i) a(i) \end{aligned}$$

先考虑如何计算 $g(i)$。先选定 $i$ 种颜色 $\binom mi$，从 $n$ 个位置 $i$ 次从挑出 $S$ 个位置，贡献为 $\dfrac{n!}{S! ^ i (n - iS)!}$，剩下的 $n - iS$ 个位置可以任意填 $m - i$ 颜色的任意一个，为 $(m - i) ^ {n - iS}$，于是 $g(i)$ 为：

$$g(i) = \binom mi \dfrac{n!(m - i) ^ {n - iS}}{S! ^ i (n - iS)!}$$

可以在 $O(m\log m)$ 时间计算。

然后考虑如何由 $g(i)$ 算出 $f(i)$，根据组合数的性质容易得到：

$$k!f(k) = \sum_{i = k} ^ m \dfrac{(-1) ^ {i - k}}{(i - k)!}i!g(i)$$

显然是一个差卷积的形式，复杂度 $O(m\log m)$，预处理阶乘及逆元为 $O(n)$。

int main()
{
	init();
	int n, m, S;
	std::cin >> n >> m >> S;
	poly w(m + 1), g(infact, infact + m + 1), f(m + 1);
	for (int i = 0; i <= m; ++ i)
		if (i & 1) adj(g[i] = -g[i]);
	for (int &x : f) scanf("%d", &x);
	for (int i = 0; i <= m; ++ i)
		if (i * S > n) w[i] = 0;
		else w[i] = (LL) C(m, i) * fact[n] % Mod * qpow(infact[S], i) % Mod
			* infact[n - i * S] % Mod * qpow(m - i, n - i * S) % Mod * fact[i] % Mod;
	std::reverse(g.begin(), g.end());
	w = w * g;
	int res = 0;
	for (int i = m; i <= 2 * m; ++ i)
		res = (res + (LL) w[i] * infact[i - m] % Mod * f[i - m]) % Mod;
	int sum = 0;
	for (int i = m; i <= 2 * m; ++ i) adj(sum += w[i] * (LL) infact[i - m] % Mod - Mod);
	std::cout << res << std::endl;
	return 0;
}