LOJ2320 「清华集训 2017」生成树计数

有难度的树上计数 + 多项式重工业。

题意：给定 $n$ 棵树，每棵树有 $a_i$ 个节点，随意连接 $n - 1$ 条边使得 $n$ 棵树联通，设连边后第 $i$ 棵树连边度数为 $d_i$，定义其权值为：

$$(\prod_{i = 1} ^ n d_i ^ m) (\sum_{i = 1} ^ n d_i ^ m)$$

求权值之和。$n\leq 3\times 10 ^ 4$，$m\leq 30$。

对 LOJ 讨论区题解的详细阐述，$m$ 可以 $10 ^ 9$。

首先不考虑 $\sum_{i = 1} ^ n d_i ^ m$ 怎么做。

直接考虑 prufer 序列，转化为序列问题，假设 $c_i$ 为 $i$ 出现的次数：

$$\begin{aligned} ans =& (n - 2)! \prod_{i = 1} ^ n \dfrac{(c_i + 1) ^ m a_i ^ {c_i + 1}} {c_i!}\\ =& (n - 2)!\prod_{i = 1} ^ n a_i \prod_{i = 1} ^ n \dfrac{(c_i + 1) ^ m a_i ^ {c_i}} {c_i!} \end{aligned}$$

考虑使用生成函数，设 $F_i(x)$ 表示 $i$ 树的生成函数，使用 EGF 即可得到：

$$F_i(x) = \sum_{j = 0} \dfrac{(j + 1) ^ m (a_ix) ^ j}{j!}$$

那么答案就是 $(n - 2)!\prod_{i = 1} ^ n a_i [x ^ {n - 2}] \prod_{i = 1} ^ n F_i(x)$。但是注意到有 $n$ 个生成函数，无法计算。

但是本题的生成函数有一个特殊的性质：$i$ 相关的只有 $(a_ix) ^ i$。这给我们一启示：按照 $a_ix$ 为一个整体先计算。

另外，所有都是乘法，不好做，乘转加使用 $\ln/\exp$，于是答案的生成函数可以写作：

$$ans = (n - 2)!\prod_{i = 1} ^ n a_i [x ^ {n - 2}]\exp\sum_{i = 1} ^ n \ln F(a_i x)$$

先计算出 $\ln F(a_i x)$，那么 $\ln$ 的和就可以写作：

$$\sum_{i = 1} ^ n \ln F(a_ix) = \sum_{j = 0} x ^ j (\sum_{i = 1} ^ n a_i ^ j) [x ^ j]\ln F(a_ix)$$

注意到有一个等幂和计算，于是可以套用 Luogu P4705 玩游戏 的套路计算即可，于是我们可以在 $O(n\log ^ 2 n)$ 的时间内得出答案的生成函数。

然后考虑加上 $\sum_{i = 1} ^ n d_i ^ m$ 如何计算。分开考虑 $i$ 的贡献，那么其实是我们把 $F_i(x)$ 的生成函数相应部分去掉，然后加上 $G_i(x)$ 的贡献。$G_i(x)$ 的生成函数比较好表达，为：

$$G_i(x) = \sum_{j = 0} \dfrac{(j + 1) ^ {2m} a_i ^ j x ^ j}{j!}$$

那么整个的生成函数可以写作：

$$(\sum_{i = 1} ^ n \dfrac{G_i(x)} {F_i(x)}) (\prod_{i = 1} ^ nF_i(x))$$

容易发现这个 $G(x)$ 也可以使用类似于 $a_ix$ 绑定一起的方式计算，那么我们可以把 $\sum_{i = 1} ^ n\dfrac{G_i(x)}{F_i(x)}$ 用刚才求等幂和的算法求出来，最后把两个生成函数乘起来，就可以得到答案了。时间复杂度 $O(n\log ^ 2 n)$。

int main()
{
	auto solve = [&](auto &self, poly &a, int l, int r) -> poly {
		if (l == r) return {1, a[l]};
		int mid = (l + r) >> 1;
		return self(self, a, l, mid) * self(self, a, mid + 1, r);
	};
	init();
	int n, m;
	std::cin >> n >> m;
	poly a(n), f(n - 1), g(n - 1);
	for (int &x : a) scanf("%d", &x);
	auto p1(solve(solve, a, 0, n - 1));
	int res = p1.back();
	for (int i = 1; i < n - 1; ++ i) res = (LL) res * i % Mod;
	for (int i = 0; i < n - 1; ++ i)
	{
		int mul = qpow(i + 1, m);
		f[i] = (LL) infact[i] * mul % Mod;
		g[i] = (LL) infact[i] * mul % Mod * mul % Mod;
	}
	g = g * Inv(f), g.resize(n - 1), p1 = Ln(p1), f = Ln(f);
	for (int i = 0; i < n; ++ i)
	{
		if (!(i & 1)) adj(p1[i] = -p1[i]);
		p1[i] = (LL) p1[i] * i % Mod;
	}
	p1[0] = n;
	for (int i = 0; i < n - 1; ++ i)
		g[i] = (LL) g[i] * p1[i] % Mod, f[i] = (LL) f[i] * p1[i] % Mod;
	res = (LL) res * (Exp(f) * g)[n - 2] % Mod;
	std::cout << res << '\n'; 
	return 0;
}