BZOJ5475 [WC2019] 数树

Description

对于两棵树 \(T_1, T_2\),定义它们的交 \(T_1\cap T_2\) 是它们的边集的交形成的森林,\(k(T_1\cap T_2)\) 表示这个森林的连通块个数,求下列三种问题之一:

  1. 给定 \(T_1, T_2\),求 \(y^{k(T_1\cap T_2)}\)
  2. 给定 \(T_1\),求 \(\sum_{T_2}y^{k(T_1\cap T_2)}\)
  3. 给定 \(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\leq k\leq 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\leq k\leq 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

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
#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\leq p_i\leq 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\leq p_i\leq k}}\prod_{i=1}^{k-2}a_{p_i}\\ =&\left(\prod_{i=1}^k a_i\right)\prod_{i=1}^{k-2}\sum_{1\leq p_i\leq k}a_{p_i}\\ =&\left(\prod_{i=1}^k a_i\right)n^{k-2}\\ \end{aligned} \]