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{i}{k}) g (i) \\ a n s = & \sum_{i = 0}^{m} f (i) a (i) \end{aligned}

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

g (i) = (\binom{m}{i}) \frac{n! (m - i)^{n - i S}}{S!^{i} (n - i S)!}

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

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

k! f (k) = \sum_{i = k}^{m} \frac{(- 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;
}