BZOJ5475 [WC2019] 数树
Description
对于两棵树 $T_1, T_2$，定义它们的交 $T_1\cap T_2$ 是它们的边集的交形成的森林，$k(T_1\cap T_2)$ 表示这个森林的连通块个数，求下列三种问题之一：
- 给定 $T_1, T_2$，求 $y^{k(T_1\cap T_2)}$；
- 给定 $T_1$，求 $\sum_{T_2}y^{k(T_1\cap T_2)}$；
- 给定 $n$，求 $\sum_{T_1}\sum_{T_2}y^{k(T_1\cap T_2)}$
其中 $n\leqslant 10^5$，上面的 $\sum$ 都是对所有可能的 $n^{n-2}$ 种树求和。
答案对 $998244353$ 取模。
Solution
Task1
直接按题意求就可以了。
Task2
如果 $y=1$，那么答案就是 $n^{n-2}$，所以下面我们假设 $y\neq 1$。
考虑令 $E_1, E_2$ 表示两棵树的边集，那么我们要求的就是 $\sum_{E_2} y^{n-|E_1\cap E_2|}$，因为有 $k$ 条边的森林恰有 $n-k$ 个连通块。我们套一个容斥原理的式子：
$$ f(S)=\sum_{T\subseteq S}\sum_{P\subseteq T}(-1)^{|T|-|P|}f(P) $$
代入 $f(S)=y^{n-|S|}, S=E_1\cap E_2$ 就可以得到
$$ \begin{aligned} &\sum_{E_2}y^{n-|E_1\cap E_2|}\\ =&\sum_{E_2}\sum_{S\subseteq E_1\cap E_2}\sum_{P\subseteq S}(-1)^{|S|-|P|}y^{n-|P|}\\ =&\sum_{S\subseteq E_1}\left(\sum_{P\subseteq S}(-1)^{|S|-|P|}y^{n-|P|}\right)g(S)\\ =&\sum_{S\subseteq E_1}y^{n-|S|}\left(\sum_{P\subseteq S}(-y)^{|S|-|P|}\right)g(S)\\ =&\sum_{S\subseteq E_1}y^{n-|S|}(1-y)^{|S|}g(S) \end{aligned} $$
其中 $g(S)$ 表示边集包含 $S$ 的树的个数。
如果 $S$ 这个边集把 $n$个点连成了一个有 $k$ 个连通块的森林，即 $n-|S|=k$，且第 $i$ 个连通块的大小为 $a_i$，那么实际上 $g(S)=n^{k-2}\prod_{i=1}^k a_i$（可以利用 Purfer 序列计算，具体见后）。
于是我们要求的东西就是
$$ \begin{aligned} &\sum_{S\subseteq E_1}y^{n-|S|}(1-y)^{|S|}g(S)\\ =&\sum_{S\subseteq E_1}y^k(1-y)^{n-k}n^{k-2}\prod_{i=1}^k a_i\\ =&\frac{(1-y)^n}{n^2}\sum_{S\subseteq E_1}\left(\frac{ny}{1-y}\right)^k\prod_{i=1}^k a_i\\ =&\frac{(1-y)^n}{n^2}\sum_{S\subseteq E_1}\prod_{i=1}^k \frac{ny}{1-y}a_i \end{aligned} $$
我们令 $k=\frac{ny}{1-y}$，那么相当于每个大小为 $a$ 的连通块会产生 $ka$ 的贡献。
考虑 DP。如果我们考虑 $i$ 所在的子树，那么 $i$ 所在的连通块有可能还会与父亲相连，所以不能立即计算它的贡献，而应该把它的大小记到状态里。
于是我们令 $f_{i,j}$ 表示在 $i$ 所在的子树里选一些边，$i$ 所在的连通块大小为 $j$ 的所有方案的权值之和（权值是指所有连通块的 $Ka$ 之积，但是 $i$ 所在的连通块是不计算的）。
那么我们要求的就是 $\sum_j kj\cdot f_{1,j}$。转移就相当于做一个树上背包，这样就有了一个 $O(n^2)$ 的算法，但这样并不能通过所有数据。
既然我们要求的只有 $k\sum_jj\cdot f_{1,j}$，我们记 $g_i=k\sum_jj\cdot f_{i,j}$。容易发现，如果我们固定 $i$ 对 $j$ 建立一个幂级数，即 $F_i(x)=\sum_{j}f_{i,j}x_j$，那么 $g_j=k\cdot F_i'(1)$。
而背包的转移就相当于
$$ F_i(x)=x\prod_j(g_j+F_j(x)) $$
其中 $j$ 遍历 $i$ 的儿子。我们发现
$$ \begin{aligned} g_i=kF_i'(1)&=k\prod_j(g_j+F_j(1))+\left(\prod_j(g_j+F_j(1))\right)\sum_j\frac{kF_j'(1)}{g_j+F_j(1)}\\ F_i(1)&=\prod_j(g_j+F_j(1)) \end{aligned} $$
第一个式子是因为（注意第二项里的 $k$ 乘到分子上了）
$$ \left(\prod_{i=1}^k F_i(x)\right)'=\sum_{i=1}^k F_i'(x)\prod_{j\neq i}F_j(x)=\left(\prod_{i=1}^kF_i(x)\right)\sum_{i=1}^k\cfrac{F_i'(x)}{F_i(x)} $$
于是我们再记一个 $t_i=F_i(1)=\sum_j f_{i,j}$，就可以得到
$$ \begin{aligned} g_i&=t_i\left(k+\sum_j\frac{g_j}{g_j+t_j}\right)\\ t_i&=\prod_j(g_j+t_j) \end{aligned} $$
于是就可以 $O(n)$ 递推了，当然你也可以随手搞一个组合意义出来什么的。
Task3
$y=1$ 同样先判掉。
类比 Task2 的做法，我们可以得到答案就是 $$ \sum_{S}y^{n-|S|}(1-y)^{|S|}g^2(S) $$ $g^2(S)$ 表示两棵树的边集都要包含 $S$。 $S$ 的枚举范围是所有森林。
展开 $g$，我们得到 $$ \frac{(1-y)^n}{n^4}\sum_{S}\prod_{i=1}^k\frac{n^2y}{1-y}a_i^2 $$ 好吧实际上就是把所有和 $g$ 有关的项都平方了。
接下来我们考虑只枚举每个连通块有多少点。如果有 $k$ 个连通块，分别有 $a_1\dots a_k$ 个点，那么组成森林的方案数就是 $\prod_i a_i^{a_i-2}$。但是 $k$ 个连通块就会算 $k!$ 次（因为实际上连通块是没有顺序的），再算上选出这些点的多重组合数 $\frac{n!}{\prod a_i!}$，于是答案就是 $$ \frac{(1-y)^nn!}{n^4}\sum_{\substack{1\leqslant k\leqslant n\\a_1+a_2+\dots+a_k=n\\a_i>0}}\frac1{k!}\prod_{i=1}^k\frac{n^2y}{(1-y)a!}a_i^{a_i} $$ 注意 $a_i^{a_i-2}$ 和 $a_i^2$ 正好合起来了。于是我们发现这其实就是一个多项式 $\exp$。 $$ \begin{aligned} &\frac{(1-y)^n}{n^4}[x^n]\sum_{1\leqslant k\leqslant n}\frac1{k!}\left(\frac{n^2y}{1-y}\sum_{a>0}\frac{a^a}{a!}x^a\right)^k\\ =&\frac{(1-y)^n}{n^4}[x^n]\exp\left(\frac{n^2y}{1-y}\sum_{a>0}\frac{a^a}{a!}x^a\right) \end{aligned} $$
于是做完了。$\exp$ 一下就好了。
Code
#include <algorithm>
#include <cstdio>
#include <cstring>
#include <set>
typedef long long LL;
const int N = 100050;
const int mod = 998244353;
const int g = 3;
LL pow_mod(LL a, int b) {
LL ans = 1;
for (; b; b >>= 1, a = a * a % mod)
if (b & 1) ans = ans * a % mod;
return ans;
}
int n, y;
namespace Task1 {
typedef std::pair<int, int> P;
std::set<P> S;
LL Work(int n, int v) {
int x, y;
for (int i = 1; i < n; ++i) {
scanf("%d%d", &x, &y);
S.insert(std::make_pair(x, y));
S.insert(std::make_pair(y, x));
}
int t = n;
for (int i = 1; i < n; ++i) {
scanf("%d%d", &x, &y);
if (S.count(std::make_pair(x, y))) --t;
}
return pow_mod(v, t);
}
};
namespace Task2 {
int pre[N], nxt[N * 2], to[N * 2], cnt;
LL K, t[N], g[N];
inline void addEdge(int x, int y) {
nxt[cnt] = pre[x];
to[pre[x] = cnt++] = y;
nxt[cnt] = pre[y];
to[pre[y] = cnt++] = x;
}
void dp(int x, int fa) {
t[x] = 1;
g[x] = K;
for (int i = pre[x]; ~i; i = nxt[i]) if (to[i] != fa) {
int y = to[i];
dp(y, x);
g[x] = (g[x] + g[y] * pow_mod(g[y] + t[y], mod - 2)) % mod;
t[x] = t[x] * (g[y] + t[y]) % mod;
}
g[x] = g[x] * t[x] % mod;
}
LL Work(int n, int y) {
if (y == 1) return pow_mod(n, n - 2);
memset(pre, -1, sizeof pre);
for (int i = 1, x, y; i < n; ++i) {
scanf("%d%d", &x, &y);
addEdge(x, y);
}
K = (LL)n * y % mod * pow_mod(1 - y, mod - 2) % mod;
dp(1, 0);
return g[1] * pow_mod(n, mod - 3) % mod * pow_mod(1 - y, n) % mod;
}
};
namespace Task3 {
const int N = ::N * 4;
LL inv[N], fac[N], ifac[N];
LL omega[N], inv_omega[N];
int rev[N], len;
void Init() {
inv[1] = 1;
for (int i = 2; i < N; ++i)
inv[i] = -(mod / i) * inv[mod % i] % mod;
fac[0] = ifac[0] = 1;
for (int i = 1; i < N; ++i) {
fac[i] = fac[i - 1] * i % mod;
ifac[i] = ifac[i - 1] * inv[i] % mod;
}
}
void InitNTT(int n) {
len = 1;
int k = 0;
while (len <= n) len <<= 1, k++;
for (int i = 1; i < len; ++i) rev[i] = (rev[i >> 1] >> 1) | ((i & 1) << (k - 1));
omega[0] = inv_omega[0] = 1;
LL v = pow_mod(g, (mod - 1) / len);
for (int i = 1; i < len; ++i)
inv_omega[len - i] = omega[i] = v * omega[i - 1] % mod;
}
void NTT(LL *A, const LL *w) {
for (int i = 0; i < len; ++i)
if (i < rev[i]) std::swap(A[i], A[rev[i]]);
for (int h = 1; h < len; h <<= 1)
for (int t = len / (h * 2), j = 0; j < len; j += h * 2) {
const LL *wn = w;
for (int i = j; i < j + h; ++i) {
LL _t1 = A[i], _t2 = A[i + h] * *wn % mod;
A[i] = (_t1 + _t2) % mod;
A[i + h] = (_t1 - _t2) % mod;
wn += t;
}
}
if (w == inv_omega)
for (int i = 0, v = -(mod - 1) / len; i < len; ++i)
A[i] = A[i] * v % mod;
}
void PolyInv(const LL *A, int n, LL *B) {
if (n == 1) { B[0] = pow_mod(A[0], mod - 2); return; }
int m = (n + 1) / 2;
PolyInv(A, m, B);
static LL tA[N], tB[N];
InitNTT(n * 2);
for (int i = 0; i < n; ++i) tA[i] = A[i];
for (int i = n; i < len; ++i) tA[i] = 0;
for (int i = 0; i < m; ++i) tB[i] = B[i];
for (int i = m; i < len; ++i) tB[i] = 0;
NTT(tA, omega); NTT(tB, omega);
for (int i = 0; i < len; ++i)
tB[i] = (2 - tA[i] * tB[i]) % mod * tB[i] % mod;
NTT(tB, inv_omega);
for (int i = 0; i < n; ++i) B[i] = tB[i];
}
void PolyLn(const LL *A, int n, LL *B) {
static LL tA[N], tB[N];
InitNTT(n * 2);
PolyInv(A, n, tA);
for (int i = n; i < len; ++i) tA[i] = 0;
for (int i = 1; i < n; ++i) tB[i - 1] = A[i] * i % mod;
for (int i = n - 1; i < len; ++i) tB[i] = 0;
NTT(tA, omega); NTT(tB, omega);
for (int i = 0; i < len; ++i)
tB[i] = tA[i] * tB[i] % mod;
NTT(tB, inv_omega);
B[0] = 0;
for (int i = 1; i < n; ++i) B[i] = tB[i - 1] * inv[i] % mod;
}
void PolyExp(const LL *A, int n, LL *B) {
if (n == 1) { B[0] = 1; return; }
int m = (n + 1) / 2;
PolyExp(A, m, B);
for (int i = m; i < n; ++i) B[i] = 0;
static LL tA[N], tB[N];
InitNTT(n * 2);
PolyLn(B, n, tA);
for (int i = 0; i < n; ++i) tA[i] = ((i == 0) + A[i] - tA[i]) % mod;
for (int i = n; i < len; ++i) tA[i] = 0;
for (int i = 0; i < m; ++i) tB[i] = B[i];
for (int i = m; i < len; ++i) tB[i] = 0;
NTT(tA, omega); NTT(tB, omega);
for (int i = 0; i < len; ++i)
tB[i] = tA[i] * tB[i] % mod;
NTT(tB, inv_omega);
for (int i = 0; i < n; ++i) B[i] = tB[i];
}
LL A[N], B[N];
LL Work(int n, int y) {
if (y == 1) return pow_mod(n, 2 * (n - 2));
Init();
LL k = (LL)n * n % mod * y % mod * pow_mod(1 - y, mod - 2) % mod;
for (int i = 1; i <= n; ++i)
A[i] = pow_mod(i, i) * ifac[i] % mod * k % mod;
PolyExp(A, n + 1, B);
return pow_mod(1 - y, n) * fac[n] % mod * pow_mod(inv[n], 4) % mod * B[n] % mod;
}
};
int main() {
freopen("tree.in", "r", stdin);
freopen("tree.out", "w", stdout);
int n, y, op;
scanf("%d%d%d", &n, &y, &op);
if (op == 0) printf("%lld\n", (Task1::Work(n, y) + mod) % mod);
if (op == 1) printf("%lld\n", (Task2::Work(n, y) + mod) % mod);
if (op == 2) printf("%lld\n", (Task3::Work(n, y) + mod) % mod);
}Prufer 序列
如果给定了一颗森林，连通块大小分别为 $a_1\dots a_k$，求其加边变成树的方案数：
可以把每个连通块看成一个点，在这些点间建出一棵树；然后如果第 $i$ 个点的度数为 $t_i$，就多乘上 $a_i^{t_i}$ 表示每条边都会在这个连通块里选一个点作为端点。
根据 Prufer 序列，这 $k$ 个点形成的树唯一对应一个长为 $k-2$ 的、每个数在 $[1,k]$ 之间的序列；且每个点的度数等于它在序列里出现的次数+1。
那么答案就是（式中 $q_i$ 表示 $p_1,p_2\dots p_{k-2}$ 中 $i$ 的出现次数）：
$$ \begin{aligned} &\sum_{\substack{p_1,p_2\dots,p_{k-2}\\1\leqslant p_i\leqslant k}}\prod_{i=1}^k a_i^{q_i+1}\\ =&\left(\prod_{i=1}^k a_i\right)\sum_{\substack{p_1,p_2\dots,p_{k-2}\\1\leqslant p_i\leqslant k}}\prod_{i=1}^{k-2}a_{p_i}\\ =&\left(\prod_{i=1}^k a_i\right)\prod_{i=1}^{k-2}\sum_{1\leqslant p_i\leqslant k}a_{p_i}\\ =&\left(\prod_{i=1}^k a_i\right)n^{k-2}\\ \end{aligned} $$