BZOJ5119 [清华集训2017] 生成树计数
Description
一个$s$个点的图，目前被$s-n$条边连成了$n$个连通块，第$i$个连通块大小为$a_i$。要求你再连$n-1$条边把它变成一棵树。对于每一棵树，若第$i$个连通块连出了$d_i$条边，其价值为
$val(T)=\left(\prod_{i=1}^nd_i^m\right)\left(\sum_{i=1}^nd_i^m\right)$
求所有生成树的价值和$\bmod998244353$。$n \leqslant 30000, m \leqslant 30$。
Note: BZOJ 上出了问题（据说评测出了问题，页面上的数据下载下来本地 50s ），可以去 LOJ 交。
Solution
首先，“连通块的大小”只是起到一个系数的作用，如果在第 $i$ 个连通块上连了 $d_i$ 条边，那么方案数要乘 $a_i^{d_i}$ 。
所以说 $$ \begin{aligned} ans&=\sum_T\left(\prod_{i=1}^na_i^{d_i}d_i^m\right)\left(\sum_{i=1}^nd_i^m\right)\\ &=\sum_{i=1}^n\sum_T\left(\prod_{i=1}^na_i^{d_i}d_i^m\right)d_i^m\\ &=\sum_{i=1}^n\sum_Ta_i^{d_i}d_i^{2m}\left(\prod_{j\neq i}a_j^{d_j}d_j^m\right) \end{aligned} $$ 推不下去了，因为“枚举所有树”太难。
我们考虑整个式子跟树的结构无关而只跟每个点的度数有关，于是可以想到 prufer 序列。
在 prufer 序列中，如果某个点的编号出现了 $k$ 次，那么它的度数即为 $k+1$ 。
于是枚举所有树可以变成枚举所有长为 $n-2$ ，元素为 $1\dots n$ 的序列；令 $k_i$ 表示序列中 $i$ 这个点出现的次数，则 $d_i=k_i+1$ 。显然如果我固定所有 $k$ 之后其方案数为 $\cfrac{(n-2)!}{\prod_i(k_i!)}$ ，所以答案即为 $$ \begin{aligned} ans&=\sum_{i=1}^n\sum_Ta_i^{d_i}d_i^{2m}\left(\prod_{j\neq i}a_i^{d_i}d_i^m\right)\\ &=\sum_{i=1}^n\sum_{k_1+k_2+\dots+k_n=n-2}\frac{(n-2)!}{\prod_i(k_i!)}a_i^{k_i+1}(k_i+1)^{2m}\left(\prod_{j\neq i}a_j^{k_j+1}(k_j+1)^m\right)\\ \end{aligned} $$ 可以发现这就是一个排列问题（枚举 $i$ 之后令第 $i$ 个点选 $k$ 个位置的价值为 $a_i^{k+1}(k+1)^{2m}$ ，其它第 $j$ 个点选 $k$ 的价值为 $a_j^{k+1}(k+1)^m$ ），于是可以想到指数型生成函数。
也即 $$ \begin{aligned} F(x)&=\sum_{i=1}^nA_i(x)\prod_{j\neq i}B_j(x)\\ A_i(x)&=\sum_ka_i^{k+1}(k+1)^{2m}\frac{x^k}{k!}\\ B_i(x)&=\sum_ka_i^{k+1}(k+1)^m\frac{x^k}{k!} \end{aligned} $$ 最终答案即为 $F$ 第 $n-2$ 项的系数乘上 $(n-2)!$ （别忘了这是指数型生成函数）。
我们发现这样复杂度是 $O(n^2\log n)$ ，无法忍受。
考虑如何简化 $A_i(x)$ 。考虑到 $m$ 相比于 $n$ 很小，根据第二类 Stirling 数的性质 $a^m=\sum_{i=0}^mS(m,i)a^{\underline i}$，我们有 $$ \begin{aligned} A_i(x)&=\sum_ka_i^{k+1}(k+1)^{2m}\frac{x^k}{k!}\\ T(x)=\int A_i(x)\mathrm{d}x&=\sum_ka_i^{k+1}(k+1)^{2m}\frac{x^{k+1}}{(k+1)!}\\ &=\sum_ka_i^kk^{2m}\frac{x^k}{k!}\\ &=\sum_{j=0}^{2m}S(2m,j)\sum_ka_i^kk^{\underline j}\frac{x^k}{k!}\\ &=\sum_{j=0}^{2m}S(2m,j)\sum_ka_i^k\frac{x^k}{(k-j)!}\\ &=\sum_{j=0}^{2m}S(2m,j)a_i^jx^je^{a_ix}\\ A_i(x)=\frac{\mathrm{d}T(x)}{\mathrm{d}x}&=\sum_{j=0}^{2m}\left[S(2m,j)ja_i^jx^{j-1}e^{a_ix}+S(2m,j)a_i^{j+1}x^je^{a_ix}\right]\\ &=e^{a_ix}\sum_{j=0}^{2m}\left[S(2m,j+1)(j+1)a_i^{j+1}+S(2m,j)a_i^{j+1}\right]x^j\\ &=e^{a_ix}\sum_{j=0}^{2m}S(2m+1,j+1)a_i^{j+1}x^j\\ \end{aligned} $$ 最后一步是由于第二类 Stirling 数的递推公式 $S(i,j)=S(i-1,j)j+S(i-1,j-1)$ 。
同理有$B_i(x)=e^{a_ix}\sum_{j=0}^{m}S(m+1,j+1)a_i^{j+1}x^j$。
所以 $$ \begin{aligned} F(x)&=\sum_{i=1}^nA_i(x)\prod_{j\neq i}B_j(x)\\ &=\sum_{i=1}^n\left(e^{a_ix}\sum_{k=0}^{2m}S(2m+1,k+1)a_i^{k+1}x^k\right)\prod_{j \neq i}\left(e^{a_jx}\sum_{k=0}^{m}S(m+1,k+1)a_j^{k+1}x^k\right)\\ &=e^{sx}\sum_{i=1}^n\left(\sum_{k=0}^{2m}S(2m+1,k+1)a_i^{k+1}x^k\right)\prod_{j \neq i}\left(\sum_{k=0}^{m}S(m+1,k+1)a_j^{k+1}x^k\right) \end{aligned} $$ 于是所有要乘起来的式子都变成了次数不超过 $2m$ 的多项式。
可以利用分治 NTT，分治过程求出 $B(x)=\prod_iB_i(x)$ 和 $A(x)=\sum_iA_i(x)\prod_{j\neq i}B_j(x)$ 。
时间复杂度 $O(nm\log^2 n)$ 。
分治时可以仅保留前 $n-2$ 项。
Code
Note: 代码中 $A(x), B(x)$ 和 Solution 中的 $A(x), B(x)$ 反了过来。
#include <algorithm>
#include <cstdio>
#include <cstring>
#include <vector>
const int N = 30050;
const int M = 35;
const int mod = 998244353;
const int g = 3;
typedef long long LL;
typedef std::vector<LL> VLL;
int n, m;
LL a[N], fac[N], ifac[N], inv[N];
inline LL pow_mod(LL x, int b) {
LL ans = 1;
for ((b += mod - 1) %= (mod - 1); b; b >>= 1, (x *= x) %= mod)
if ((b & 1) != 0) (ans *= x) %= mod;
return ans;
}
void NTT(LL *A, int len, int opt) {
for (int i = 1, j = 0; i < len; ++i) {
int k = len;
do j ^= (k >>= 1); while ((j & k) == 0);
if (i < j) std::swap(A[i], A[j]);
}
for (int h = 2; h <= len; h <<= 1) {
LL wn = pow_mod(g, (mod - 1) / h * opt);
for (int j = 0; j < len; j += h) {
LL w = 1LL;
for (int i = j; i < j + (h >> 1); ++i) {
LL _tmp1 = A[i], _tmp2 = A[i + (h >> 1)] * w % mod;
A[i] = (_tmp1 + _tmp2) % mod;
A[i + (h >> 1)] = (_tmp1 - _tmp2) % mod;
(w *= wn) %= mod;
}
}
}
if (opt == -1)
for (int i = 0; i < len; ++i)
(A[i] *= -(mod - 1) / len) %= mod;
}
LL S[2 * M][2 * M];
struct PVLL{
VLL A, B;
PVLL() : A(0), B(0) {}
};
inline void Copy(const VLL &x, LL *A, int len) {
for (int i = 0; i < len; ++i) A[i] = (i < x.size() ? x[i] : 0);
}
PVLL operator*(const PVLL &x, const PVLL &y) {
static LL T1[N * M * 4], T2[N * M * 4], T3[N * M * 4];
PVLL ans;
int len = 1,
ansl = std::max(x.A.size() + y.B.size(), x.B.size() + y.A.size());
while (len < ansl - 1) len <<= 1;
Copy(x.A, T1, len);
NTT(T1, len, 1);
Copy(y.A, T2, len);
NTT(T2, len, 1);
for (int i = 0; i < len; ++i) T3[i] = T1[i] * T2[i] % mod;
NTT(T3, len, -1);
ans.A.resize(std::min<int>(n, x.A.size() + y.A.size() - 1));
for (int i = 0; i < std::min<int>(n, x.A.size() + y.A.size() - 1); ++i)
(ans.A[i] = T3[i]) %= mod;
Copy(x.B, T3, len);
NTT(T3, len, 1);
for (int i = 0; i < len; ++i) (T3[i] *= T2[i]) %= mod;
NTT(T3, len, -1);
ans.B.resize(std::min(n, ansl - 1));
for (int i = 0; i < std::min(n, ansl - 1); ++i)
ans.B[i] = T3[i];
Copy(y.B, T3, len);
NTT(T3, len, 1);
for (int i = 0; i < len; ++i) (T3[i] *= T1[i]) %= mod;
NTT(T3, len, -1);
for (int i = 0; i < std::min(n, ansl - 1); ++i)
(ans.B[i] += T3[i]) %= mod;
return ans;
}
PVLL Solve(int l, int r) {
if (l == r - 1) {
PVLL ans;
ans.A.resize(m + 1);
ans.B.resize(2 * m + 1);
LL pa = a[l];
for (int i = 0; i <= 2 * m; ++i, (pa *= a[l]) %= mod) {
if (i <= m) ans.A[i] = (LL)pa * S[m + 1][i + 1] % mod;
ans.B[i] = (LL)pa * S[2 * m + 1][i + 1] % mod;
}
return ans;
}
int mid = (l + r + 1) >> 1;
return Solve(l, mid) * Solve(mid, r);
}
LL ps[N];
int main() {
scanf("%d%d", &n, &m);
fac[0] = fac[1] = ifac[0] = ifac[1] = inv[1] = 1;
for (int i = 2; i <= n; ++i) {
fac[i] = fac[i - 1] * i % mod;
inv[i] = -(mod / i) * inv[mod % i] % mod;
ifac[i] = ifac[i - 1] * inv[i] % mod;
}
S[0][0] = 1;
for (int i = 1; i <= 2 * m + 1; ++i)
for (int j = 1; j <= 2 * m + 1; ++j)
S[i][j] = (S[i - 1][j - 1] + S[i - 1][j] * j) % mod;
int s = 0;
for (int i = 0; i < n; ++i) {
scanf("%lld", &a[i]);
(s += a[i]) %= mod;
}
PVLL res = Solve(0, n);
LL ans = 0;
ps[0] = 1;
for (int i = 1; i <= n - 2; ++i) ps[i] = ps[i - 1] * s % mod;
for (int i = 0; i <= n - 2; ++i)
(ans += res.B[i] * ps[n - 2 - i] % mod * ifac[n - 2 - i] % mod) %= mod;
printf("%lld\n", (ans * fac[n - 2] % mod + mod) % mod);
return 0;
}