【LOJ #2983】「WC2019」数树
Description
题目链接：
本题包含三个问题：
- 问题 0：已知两棵 \(n\) 个节点的树的形态（两棵树的节点标号均为 \(1\) 至 \(n\)），其中第一棵树是红树，第二棵树是蓝树。要给予每个节点一个 \([1, y]\) 中的整数，使得对于任意两个节点 \(p, q\)，如果存在一条路径 \((a_1 = p, a_2, \cdots , a_m = q)\) 同时属于这两棵树，则 \(p, q\) 必须被给予相同的数。求给予数的方案数。
- 问题 1：已知蓝树，对于红树的所有 \(n^{n−2}\) 种选择方案，求问题 0 的答案之和。
- 问题 2：对于蓝树的所有 \(n^{n−2}\) 种选择方案，求问题 1 的答案之和。
输出答案对 \(998, 244, 353\) 取模的结果。
\(n \leq 10^5\)
时空限制：\(\texttt{4s/512MB}\)
Solution
注意 \(y=1\) 的情况会带来一些困扰，提前判掉即可。
问题 0
根据题意，由两棵树边集的交集构成的同一个连通块内的点均要取同一个颜色，森林的连通块数等于点数减边数，因此求出边集交集的大小 \(k\)，答案就是 \(y^{n-k}\)。
可以用 \(\text{std::map}\) 实现这个过程。
问题 1
暴力矩阵树
假设每个生成树的权值为 \(y^{n-k}\)，\(k\) 表示这棵生成树和给定的树的边集交集大小，那么要求的就是所有生成树的权值之和。不妨将 \(y^n\) 提取出来，那么每条原树的边的边权都是 \(y^{-1}\)，不是原树的边的边权就是 \(1\)。
直接套用矩阵树定理即可。时间复杂度 \(\mathcal O(n^3)\)。
听说这个矩阵比较有性质可以优化到 \(\mathcal O(n)\)，因为我太菜无法理解题解的做法，所以这种做法就到此为止了。
推式子
不妨将第一棵树的边集叫做 \(E_1\)，第二棵树的边集叫做 \(E_2\)，那么就有下面这个式子成立
\[\begin{aligned}
\mathrm{ans}&=\sum_{E_2}y^{n-|E_1\cap E_2|} \\
&=\sum_{S\subseteq E_1}y^{n-|S|}\sum_{E_2}[S=E_1\cap E_2]
\end{aligned}
\]
不妨设
\[f(S)=\sum_{E_2}[S=E_1\cap E_2]
\]
那么 \(f(S)\) 显然不太好求，不妨考虑一下 \(f(S)\) 的实际含义，\(f(S)\) 就表示将强制 \(E_2\) 首先要选边集 \(S\) 的边，并且强制 \(E_2\) 不能选 \(E_1\setminus S\) 的边。
前者比较好处理，但是后者不太好处理。我们考虑先去掉后面那个条件，就想到设
\[g(S)=\sum_{E_2}[S\subseteq E_1\cap E_2]
\]
显然有
\[g(S)=\sum_{S\subseteq T}f(T)
\]
容斥一下就是
\[f(S)=\sum_{S\subseteq T}(-1)^{|T|-|S|}g(T)
\]
所以
\[\begin{aligned}
\mathrm{ans}&=\sum_{S\subseteq E_1}y^{n-|S|}f(S)\\
&=\sum_{S\subseteq E_1}y^{n-|S|}\sum_{S\subseteq T\subseteq E_1}(-1)^{|T|-|S|}g(T)\\
&=\sum_{T\subseteq E_1}g(T)\sum_{S\subseteq T}y^{n-|S|}(-1)^{|T|-|S|}\\
&=y^n\sum_{T\subseteq E_1}g(T)\sum_{i=0}^{|T|}\binom{|T|}{i}y^{-i}(-1)^{|T|-i}\\
&=y^n\sum_{T\subseteq E_1}g(T)(y^{-1}-1)^{|T|}
\end{aligned}
\]
如何求 \(g(T)\)
现在先考虑 \(g(T)\) 是怎么求的，从我们为什么要定义 \(g(T)\) 那里我们已经知道了 \(g(T)\) 的具体含义：强制 \(E_2\) 一定要选 \(T\) 的边，\(E_2\) 的方案数。
这相当于，我们已经确定了一棵有标号无根树中的若干条边，考虑剩下的边的选取方案。不妨考虑这个边集 \(T\)，已经将当前的点集并成 \(k\) 个连通块，第 \(i\) 个连通块的大小为 \(a_i\)。我们可以将每个连通块看成一个点，那么每条连接两个连通块 \(a_i,a_j\) 之间的边都可以有 \(a_i\cdot a_j\) 种选择。接下去的推导就可以用 prufer 序列或者矩阵树定理了。
prufer 序列推导
设 \(p_1,p_2,\cdots,p_{k-2}\) 表示 prufer 序列，\(q_i\) 表示 \(i\) 在 prufer 序列中的出现次数。因为每个点的度数为 \(q_i+1\)，那么第 \(i\) 个连通块的贡献就是 \(a_i^{q_i+1}\)。具体地
\[\begin{aligned}
g(T)&=\sum_{\{p_{k-2}\}}\prod_{i=1}^ka_i^{q_i+1}\\
&=\prod_{i=1}^ka_i\sum_{\{p_{k-2}\}}\prod_{i=1}^ka_i^{q_i}\\
&=\prod_{i=1}^ka_i\sum_{\{p_{k-2}\}}\prod_{i=1}^{k-2}a_{p_i}\\
&=\prod_{i=1}^ka_i\prod_{i=1}^{k-2}\sum_{j=1}^ka_{j}\\
&=n^{k-2}\prod_{i=1}^ka_i
\end{aligned}
\]
矩阵树定理推导
去掉最后一行最后一列的基尔霍夫矩阵的行列式是
\[\begin{vmatrix}
a_1(n-a_1) & -a_1a_2 & \cdots & -a_1a_{k-1} \\
-a_1a_2 & a_2(n-a_2) & \cdots & -a_2a_{k-1} \\
\cdots & \cdots & \cdots & \cdots \\
-a_1a_{k-1}& -a_2a_{k-1}& \cdots & a_{k-1}(n-a_{k-1})
\end{vmatrix}
\]
把每一行的 \(a_i\) 提出来就是
\[\prod_{i=1}^{k-1}a_i
\begin{vmatrix}
n-a_1 & -a_2 & \cdots & -a_{k-1} \\
-a_1 & n-a_2 & \cdots & -a_{k-1} \\
\cdots & \cdots & \cdots & \cdots \\
-a_1 & -a_2 & \cdots & n-a_{k-1}
\end{vmatrix}
\]
前面 \(k-2\) 行全部减去最后一行
\[\prod_{i=1}^{k-1}a_i
\begin{vmatrix}
n & 0 & \cdots & 0 & -n \\
0 & n & \cdots & 0 & -n\\
\cdots & \cdots & \cdots & \cdots & \cdots \\
0 & 0 & \cdots & n & -n \\
-a_1 & -a_2& \cdots & -a_{k-2} & n-a_{k-1}
\end{vmatrix}
\]
前面 \(k-2\) 列全部加到最后一列上
\[\prod_{i=1}^{k-1}a_i
\begin{vmatrix}
n & 0 & \cdots & 0 & 0 \\
0 & n & \cdots & 0 & 0\\
\cdots & \cdots & \cdots & \cdots & \cdots \\
0 & 0 & \cdots & n & 0 \\
-a_1 & -a_2& \cdots & -a_{k-2} & a_k
\end{vmatrix}
\]
于是
\[g(T)=\prod_{i=1}^{k-1}a_i
\begin{vmatrix}
n & 0 & \cdots & 0 & 0 \\
0 & n & \cdots & 0 & 0\\
\cdots & \cdots & \cdots & \cdots & \cdots \\
0 & 0 & \cdots & n & 0 \\
-a_1 & -a_2& \cdots & -a_{k-2} & a_k
\end{vmatrix}=n^{k-2}\prod_{i=1}^{k}a_i
\]
继续推式子
推导出这个以后我们继续考虑推式子。我们设 \(R(T)\) 表示固定了边集 \(T\) 后得到的所有连通块的大小之积。
\[\begin{aligned}
\mathrm{ans}&=y^n\sum_{T\subseteq E_1}g(T)(y^{-1}-1)^{|T|}\\
&=y^n\sum_{T\subseteq E_1}n^{n-|T|-2}R(T)(y^{-1}-1)^{|T|}\\
&=\frac{y^nn^n}{n^2}\sum_{T\subseteq E_1}R(T)\left(\frac{y^{-1}-1}{n}\right)^{|T|}\\
\end{aligned}
\]
我们不妨设 \(a(T)_i\) 表示 \(T\) 得到的第 \(i\) 个连通块大小，那么有
\[\begin{aligned}
\mathrm{ans}&=\frac{y^nn^n}{n^2}\sum_{T\subseteq E_1}\left(\frac{y^{-1}-1}{n}\right)^{|T|}\prod_{i=1}^{n-|T|}a(T)_i\\
&=\frac{(1-y)^n}{n^2}\sum_{T\subseteq E_1}\prod_{i=1}^{n-|T|}a(T)_i\left(\frac{n}{y^{-1}-1}\right)
\end{aligned}
\]
树形 DP
不妨设 \(K=\frac{n}{y^{-1}-1}\)，现在只需要考虑计算
\[\sum_{T\subseteq E_1}\prod_{i=1}^{n-|T|}K\cdot a(T)_i
\]
我们发现边集与每种划分连通块的方案形成双射。现在可以考虑将枚举边集看成划分连通块。
相当于对于每种划分连通块的方案，每个大小为 \(s\) 的连通块的权值是 \(Ks\)，每个方案的权值是所有连通块的权值之积，对所有方案的权值求和。这个显然是可以把当前点的连通块大小记到状态里 \(\mathcal O(n^2)\) 树形 DP 做的，但是不足以通过本题。
考虑这个式子的组合意义，相当于将整棵树划分成若干连通块，然后在每个连通块内选一个点，并乘上 \(K\) 的权值。那么考虑到这里我们就可以 \(\mathcal O(n)\) 树形 DP 了。
我们设 \(f(u,0/1)\) 表示考虑到 \(u\) 的子树，并且 \(u\) 所在的连通块是否选择了点的所有方案权值之和。
转移可以看代码。
inline void dfs(int u, int pre)
{
f[u][0] = 1;
f[u][1] = K;
foredge(u)
if (v != pre)
{
dfs(v, u);
int f0 = (1LL * f[u][0] * f[v][0] + 1LL * f[u][0] * f[v][1]) % mod;
int f1 = (1LL * f[u][1] * f[v][1] + 1LL * f[u][1] * f[v][0]) % mod;
add(f1, 1LL * f[u][0] * f[v][1] % mod);
f[u][0] = f0;
f[u][1] = f1;
}
}
问题 2
推式子是类似的，不过这时候就没有限制 \(T\subseteq E_1\) 了，\(T\) 可以选取任意无环无重边边集。
\[\mathrm{ans}=y^n\sum_{T}g(T)^2(y^{-1}-1)^{|T|}
\]
代入 \(g(T)\) 的表达式我们有
\[\begin{aligned}
\mathrm{ans}&=(1-y)^n\sum_{T}R(T)^2n^{2(n-|T|-2)}\left(\frac{1}{y^{-1}-1}\right)^{n-|T|}\\
&=\frac{(1-y)^n}{n^4}\sum_{T}R(T)^2\left(\frac{n^2}{y^{-1}-1}\right)^{n-|T|}\\
&=\frac{(1-y)^n}{n^4}\sum_{T}\prod_{i=1}^{n-|T|}a(T)_i^2\left(\frac{n^2}{y^{-1}-1}\right)\\
\end{aligned}
\]
与问题 1 不同的地方体现在，问题 2 没有给出一棵树，因此每种连通块划分方案对应着多个边集，需要在每个连通块内部任意连一棵树出来。而且有标号使得划分连通块大小后还需要给每个连通块分配标号。
因此这涉及到有标号组合对象的拼接，我们考虑指数生成函数，设 \(F(x)\) 表示每个连通块的指数生成函数，那么
\[\begin{aligned}
F(x)&=\sum_{i=1}^{+\infty}\frac{i^2\cdot i^{i-2}\cdot n^2}{y^{-1}-1} \cdot\frac{x^i}{i!}\\
&=\sum_{i=1}^{+\infty}\frac{i^i\cdot n^2}{y^{-1}-1} \cdot\frac{x^i}{i!}
\end{aligned}
\]
那么就有
\[\begin{aligned}
\mathrm{ans}&=\frac{(1-y)^nn!}{n^4}[x^n]\sum_{k=1}^n\frac{F(x)^k}{k!}\\
&=\frac{(1-y)^nn!}{n^4}[x^n]e^{F(x)}
\end{aligned}
\]
需要除以 \(k!\) 是因为连通块是无标号的，然后套个多项式 exp 的板子就行了。
时间复杂度 \(\mathcal O(n \log n)\)。
#include <bits/stdc++.h>
template <class T>
inline void read(T &x)
{
static char ch;
while (!isdigit(ch = getchar()));
x = ch - '0';
while (isdigit(ch = getchar()))
x = x * 10 + ch - '0';
}
typedef long long s64;
const int mod = 998244353;
inline void add(int &x, const int &y)
{
x += y;
if (x >= mod)
x -= mod;
}
inline void dec(int &x, const int &y)
{
x -= y;
if (x < 0)
x += mod;
}
inline int qpow(int x, int y)
{
int res = 1;
for (; y; y >>= 1, x = 1LL * x * x % mod)
if (y & 1)
res = 1LL * res * x % mod;
return res;
}
const int MaxN = 1e6 + 5;
int n, y, opt;
namespace task0
{
std::map<int, bool> bo[MaxN];
inline void answer()
{
for (int i = 1; i < n; ++i)
{
int u, v;
read(u), read(v);
if (u > v)
std::swap(u, v);
bo[u][v] = true;
}
int res = n;
for (int i = 1; i < n; ++i)
{
int u, v;
read(u), read(v);
if (u > v)
std::swap(u, v);
res -= bo[u][v];
}
std::cout << qpow(y, res) << '\n';
}
}
namespace task1
{
int K;
int f[MaxN][2];
int ect, adj[MaxN];
int nxt[MaxN << 1], to[MaxN << 1];
#define foredge(u) for (int e = adj[u], v; v = to[e], e; e = nxt[e])
inline void addEdge(int u, int v)
{
nxt[++ect] = adj[u];
adj[u] = ect;
to[ect] = v;
}
inline void dfs(int u, int pre)
{
f[u][0] = 1;
f[u][1] = K;
foredge(u)
if (v != pre)
{
dfs(v, u);
int f0 = (1LL * f[u][0] * f[v][0] + 1LL * f[u][0] * f[v][1]) % mod;
int f1 = (1LL * f[u][1] * f[v][1] + 1LL * f[u][1] * f[v][0]) % mod;
add(f1, 1LL * f[u][0] * f[v][1] % mod);
f[u][0] = f0;
f[u][1] = f1;
}
}
inline void answer()
{
for (int i = 1; i < n; ++i)
{
int u, v;
read(u), read(v);
addEdge(u, v);
addEdge(v, u);
}
K = 1LL * n * qpow(qpow(y, mod - 2) - 1, mod - 2) % mod;
dfs(1, 0);
int d = 1LL * qpow(mod - y + 1, n) * qpow(1LL * n * n % mod, mod - 2) % mod;
std::cout << 1LL * f[1][1] * d % mod << '\n';
}
}
namespace polynomial
{
int inv[MaxN];
int rev[MaxN], P, L;
inline void inv_init()
{
inv[1] = 1;
for (int i = 2; i < MaxN; ++i)
inv[i] = 1LL * inv[mod % i] * (mod - mod / i) % mod;
}
inline void DFT_init(int n)
{
P = 0, L = 1;
while (L < n)
{
L <<= 1;
++P;
}
for (int i = 1; i < L; ++i)
rev[i] = (rev[i >> 1] >> 1) | ((i & 1) << (P - 1));
}
inline void resize(int *a, int cur_n, int n)
{
for (int i = n; i < cur_n; ++i)
a[i] = 0;
}
inline void DFT(int *a, int n, int opt)
{
for (int i = 0; i < n; ++i)
if (i < rev[i])
std::swap(a[i], a[rev[i]]);
int g = opt == 1 ? 3 : (mod + 1) / 3;
for (int k = 1; k < n; k <<= 1)
{
int omega = qpow(g, (mod - 1) / (k << 1));
for (int i = 0; i < n; i += k << 1)
{
int x = 1;
for (int j = 0; j < k; ++j)
{
int u = a[i + j];
int v = 1LL * x * a[i + j + k] % mod;
add(a[i + j] = u, v);
dec(a[i + j + k] = u, v);
x = 1LL * x * omega % mod;
}
}
}
if (opt == -1)
{
for (int i = 0; i < n; ++i)
a[i] = 1LL * a[i] * inv[n] % mod;
}
}
inline void dot_mul(int *a, int *b, int *c)
{
for (int i = 0; i < L; ++i)
c[i] = 1LL * a[i] * b[i] % mod;
}
inline void mul(int *a, int *b, int *c, int n, int m)
{
DFT_init(n + m - 1);
static int ta[MaxN], tb[MaxN];
for (int i = 0; i < L; ++i) ta[i] = tb[i] = 0;
for (int i = 0; i < n; ++i) ta[i] = a[i];
for (int i = 0; i < m; ++i) tb[i] = b[i];
DFT(ta, L, 1), DFT(tb, L, 1);
dot_mul(ta, tb, c), DFT(c, L, -1), resize(c, L, n + m - 1);
}
inline void inverse(int *a, int *b, int n)
{
static int ta[MaxN], tb[MaxN], m;
for (int i = 0; i < (n << 2); ++i)
ta[i] = tb[i] = b[i] = 0;
b[0] = qpow(a[0], mod - 2);
for (m = 1; m < n; m <<= 1)
{
DFT_init(m << 2);
for (int i = 0; i < (m << 1); ++i)
ta[i] = a[i], tb[i] = b[i];
DFT(ta, L, 1), DFT(tb, L, 1);
for (int i = 0; i < L; ++i)
{
int d = 1LL * ta[i] * tb[i] % mod;
b[i] = 1LL * tb[i] * (2 + mod - d) % mod;
}
DFT(b, L, -1);
resize(b, L, m << 1);
}
resize(b, m, n);
}
inline void derivative(int *a, int *b, int n)
{
for (int i = 0; i < n - 1; ++i)
b[i] = 1LL * (i + 1) * a[i + 1] % mod;
b[n - 1] = 0;
}
inline void anti_derivative(int *a, int *b, int n)
{
for (int i = n; i >= 1; --i)
b[i] = 1LL * inv[i] * a[i - 1] % mod;
b[0] = 0;
}
inline void ln(int *a, int *b, int n)
{
static int ta[MaxN], tb[MaxN];
derivative(a, ta, n);
inverse(a, tb, n);
mul(tb, ta, b, n, n - 1);
anti_derivative(b, b, n), resize(b, L, n);
}
inline void exp(int *a, int *b, int n)
{
static int tb[MaxN], m;
for (int i = 0; i < (n << 2); ++i)
tb[i] = b[i] = 0;
b[0] = 1;
for (m = 1; m < n; m <<= 1)
{
ln(b, tb, m << 1);
tb[0] = (a[0] + mod + 1 - tb[0]) % mod;
for (int i = 1; i < (m << 1); ++i)
tb[i] = (mod - tb[i] + a[i]) % mod;
mul(b, tb, b, m, m << 1);
resize(b, L, m << 1);
}
resize(b, m, n);
}
}
namespace task2
{
int fac[MaxN], fac_inv[MaxN];
inline void fac_init(int n)
{
fac[0] = 1;
for (int i = 1; i <= n; ++i)
fac[i] = 1LL * fac[i - 1] * i % mod;
fac_inv[n] = qpow(fac[n], mod - 2);
for (int i = n - 1; i >= 0; --i)
fac_inv[i] = 1LL * fac_inv[i + 1] * (i + 1) % mod;
}
inline void answer()
{
polynomial::inv_init();
fac_init(n);
static int f[MaxN], g[MaxN];
int d = 1LL * n * n % mod * qpow(qpow(y, mod - 2) - 1, mod - 2) % mod;
for (int i = 1; i <= n; ++i)
f[i] = 1LL * d * qpow(i, i) % mod * fac_inv[i] % mod;
polynomial::exp(f, g, n + 1);
int t = 1LL * fac[n] * qpow(mod - y + 1, n) % mod;
t = 1LL * t * qpow(n, (4LL * (mod - 2)) % (mod - 1)) % mod;
std::cout << 1LL * t * g[n] % mod << '\n';
}
}
int main()
{
read(n), read(y), read(opt);
if (opt == 0)
task0::answer();
else
{
if (y == 1)
std::cout << qpow(n, opt * (n - 2)) << '\n';
else
{
if (opt == 1)
task1::answer();
else
task2::answer();
}
}
return 0;
}
