[LOJ#2983][WC2019]数树

Address

• 洛谷 P5206
• LOJ #2983

Solution

• 这是一道计数 (shen) 好 (xian) 题，综合性很强 ，虽然是三合一
• 鉴于此题用到的算法较多，我就不在标题里列出了，此题用到的算法及技巧如下
• 容斥原理
• prufer 序列计数
• 组合数学
• 树形 DP
• 多项式 exp

Task 1 ： $op=0$

• 首先分析下问题，很容易得出，如果设 $T_0$ 表示一个 $n$ 个点的森林，边集为两棵树 $T_1$ 和 $T_2$ 的边交集，且 $T_0$ 包含 $cnt$ 个连通块，则答案为
• $$y^{cnt}$$
• 利用 map 等很多技巧都可以求两棵树的边交集
• $O(n\log n)$

Task 2 ： $op=1$

前置知识：已经确定一些边的生成树计数

• 这个问题大致为：有 $m$ 个连通块，第 $i$ 个连通块包含 $a_i$ 个点， $n={\sum }_{i=1}^m{a}_i$
• 求用 $m-1$ 条边连接 $m$ 个连通块的方案数
• 考虑一种对树进行计数的工具 prufer 序列
• 先考虑，如果对于所有的 $i$ 都有 $a_i=1$ ，那么长度为 $m-2$ 的 prufer 序列每个元素都可以取 $m$ 个点中的任一个，方案数 $m^{m-2}$
• 此外， prufer 序列的另一个性质：度数为 $d$ 的点在序列中出现 $d-1$ 次
• 回到 $a_i$ 不一定等于 $1$ 的情况，设第 $i$ 的点的度数为 $d_i$
• 那么第 $i$ 个连通块会让方案数额外乘上 $a_i^{d_i}$
• 我们可以把过程看成如下
• 初始时设 $res={\prod }_{i=1}^m{a}_i$
• 然后考虑 prufer 序列的第 $i$ 个元素 $p_i$
• 如果 $p_i$ 取 $u$ ，那么会为方案数乘上 $a_u$ ，而 $p_i$ 可以取 $m$ 个连通块中的任一，故将 $res$ 乘上 $n$
• 注意到处理 prufer 的过程中，第 $i$ 个连通块贡献了 $d_i-1$ 次 $a_i$，再乘上初始的 ${\prod }_{i=1}^m{a}_i$ 恰好是贡献 $d_i$ 次
• 所以方案数为
• $$res=n^{m-2}\prod _{i=1}^m{a}_i$$

回到问题

• 我们知道，对于森林，有：边数 $+$ 连通块个数 $=$ 点数
• 设 $C(S)$ 表示树 $T_2$ 有多少种取法使得 $T_1$ 与 $T_2$ 的边交集恰好为 $S$
• 那么答案显然为
• $$\sum _{S}C(S)\times y^{n-|S|}$$
• 考虑利用容斥化一下 $C(S)$
• 设 $D(S)$ 表示树 $T_2$ 有多少种取法使得边集 $S$ 是 $T_1$ 与 $T_2$ 的边交集的子集
• 那么 $D(S)$ 显然比 $C(S)$ 容易直接求。这就相当于 $T_1$ 和 $T_2$ 都必须包含边集 $S$
• 于是 $D(S)$ 就相当于有 $|S|$ 条边已经确定的情况下，在剩下的边中选 $n-1-|S|$ 条，连通所有点的方案数，可以用上面介绍的 prufer 序列计数求得
• 于是求 $C(S)$ 我们可以容斥
• $$C(S)=\sum _{S\subset T}(-1{)}^{|T|-|S|}D(T)$$
• 代入答案的式子
• $$answer=y^n\sum _S(y^{-1}{)}^{|S|}C(S)$$
• $$=y^n\sum _S(y^{-1}{)}^{|S|}\sum _{S\subset T}(-1{)}^{|T|-|S|}D(T)$$
• 把 $T$ 集合提到外层枚举
• $$=y^n\sum _{T}D(T)\sum _{S\subset T}(y^{-1}{)}^{|S|}(-1{)}^{|T|-|S|}$$
• 发现在枚举 $S\subset T$ 时， $|S|$ 相同的集合 $S$ 本质上没有区别
• $$=y^n\sum _{T}D(T)\sum _{i=0}^{|T|}\left(\genfrac{}{}{0ex}{}{|T|}{i}\right)(y^{-1}{)}^i(-1{)}^{|T|-i}$$
• 由二项式定理得
• $$=y^n\sum _{T}D(T)(y^{-1}-1{)}^{|T|}$$
• 注意这里的 $T$ 只能取树 $T_1$ 的边子集
• 发现当 $y=1$ 时会对我们的处理造成不便，直接特判掉： $y=1$ 时答案为 $n^{n-2}$
• 然后继续转化上式，设 $R(S)$ 表示 $n$ 个点，加入树 $T_1$ 中除边集 $S$ 之外的所有边之后，每个连通块大小之积，根据上面介绍的 prufer 序列计数，可以得到
• $$=y^n\sum _S{n}^{|S|-1}R(S)(y^{-1}-1{)}^{n-(|S|+1)}$$
• $$=\frac{(y(y^{-1}-1){)}^n}{n^2}\sum _{S}R(S)(n(y^{-1}-1{)}^{-1}{)}^{|S|+1}$$
• $\frac{(y(y^{-1}-1){)}^n}{n^2}$ 这一部分可以在算好后面的部分之后再去乘上。我们设 $Y=n(y^{-1}-1{)}^{-1}$ ，那么 ${\sum }_{S}R(S)Y^{|S|+1}$ 就相当于把树 $T_1$ 分成几个连通块，分成 $i$ 个连通块的贡献为 $Y^i$ 乘上所有连通块的大小之积
• 这显然可以用树形 DP 解决
• 发现由于需要考虑所有连通块的大小之积，所以这个 DP 是 $O(n^2)$ 的
• 我们不妨寻找一个 $R(S)Y^{|S|+1}$ 的组合意义
• 这相当于在 $|S|+1$ 个连通块中，每个连通块内各选一点，选一个点会乘上 $Y$ 的贡献
• 于是 ${\sum }_{S}R(S)Y^{|S|+1}$ 就相当于选出一些点，若选出了 $i$ 个点，那么贡献为 $Y^i$ 乘上选出 $i-1$ 条边切掉使得这 $i$ 个点两两不连通的方案数
• 于是我们可以开始 DP 了
• $f[u]$ 表示 $u$ 的子树内选出一些点，使得这些点归属不同的连通块
• $g[u]$ 表示 $u$ 的子树内选出一些点（不能选 $u$ ），使得 $u$ 归属一个连通块，选出的点各自位于一个连通块
• 如果 $u$ 的子树只有一个点则 $f[u]=Y$ ， $g[u]=1$
• 使用背包合并的方式转移，如果 $u$ 处理到子树 $v$ 之前的 DP 数组为 $f^{\prime }$ 和 $g^{\prime }$ ，则
• $$f[u]=f^{\prime }[u]\times f[v]+g^{\prime }[u]\times f[v]+f^{\prime }[u]\times g[v]$$
• $$g[u]=g^{\prime }[u]\times f[v]+g^{\prime }[u]\times g[v]$$
• 即转移时分 $(u,v)$ 切掉与不切掉进行讨论
• 于是答案为
• $$\frac{(y(y^{-1}-1){)}^n}{n^2}f[Root]$$
• $O(n)$

Task 3 ： $op=2$

前置知识：指数生成函数与多项式 EXP

• 定义组合对象 $A$ 的指数生成函数 $F(x)$ 为
• $$F(x)=\sum _{i\ge 0}\frac{A_i}{i!}$$
• 可以得出，如果组合对象 $A$ 和 $B$ 的指数生成函数分别为 $F(x)$ 和 $G(x)$ ，那么将这两个组合对象进行带标号拼接得到 $C$ ，则 $C$ 的指数生成函数 $H(x)$ 有
• $$[x^n]H(x)=\frac{C_n}{n!}=\sum _{i+j=n}\left(\genfrac{}{}{0ex}{}{n}{i}\right)A_i{B}_j(n!{)}^{-1}$$
• $$=\sum _{i+j=n}[x^i]F(x)[x^j]G(x)$$
• 即 $H(x)=F(x)G(x)$
• 然后引入多项式 EXP
• 设组合对象 $A$ 的指数生成函数 $F(x)$ （ $A_0=0$ ），有
• $$e^{F(x)}=\sum _{i\ge 0}\frac{F^i(x)}{i!}$$
• 显然 $F^i(x)$ 表示 $i$ 个组合对象 $A$ 带标号拼接
• $\frac{F^i(x)}{i!}$ 表示 $i$ 个组合对象 $A$ 不考虑彼此顺序的拼接
• 于是 $e^{F(x)}$ 表示一些组合对象 $A$ 进行不考虑彼此顺序的拼接
• 以上概念可能有些含糊，所以我们引入一个例子
• 如，无向简单图（不含重边和自环）是由若干无向简单连通图构成的
• 连通块之间没有顺序
• 所以若无向简单图个数的指数生成函数为 $G(x)$ ，无向简单连通图个数的指数生成函数为 $F(x)$ ，则
• $$G(x)=e^{F(x)}$$
• 显然 $G(x)$ 是个形式幂级数，我们如何求得 $G(x)\mathrm{m}\mathrm{o}\mathrm{d}{x}^n$ 呢？
• 由牛顿迭代得
• $$G_{k+1}(x)=G_k(x)(1-\ln G_k(x)+F(x))$$
• 其中 $G_k(x)=G(x)\mathrm{m}\mathrm{o}\mathrm{d}{x}^{2^k},G_0(x)=1$
• 将 $k$ 从 $1$ 迭代到 $\ge n$ 为止即可。复杂度 $O(n\log n)$

回到问题

• 特判掉 $y=1$ 时答案为 $n^{2(n-2)}$
• 前面的容斥推导基本一样，同样地
• $$answer=y^n\sum _{T}D(T)(y^{-1}-1{)}^{|T|}$$
• 注意，这时候 $D(T)$ 表示钦定边集 $T$ 时 $n$ 个点生成树个数的平方（因为要确定两棵树），这里的 $T$ 可以取任意不成环的边集
• 易得 $D(T)$ 等于 $n$ 的 $2(n-|T|-2)$ 次幂，乘上所有连通块大小之积的平方，但如果我们枚举的是连通块分属情况而不是边集 $T$ ，我们还需要考虑每个连通块长什么样
• 我们尝试把枚举 $T$ 改成枚举边集 $T$ 的连通块分属情况，则
• $$answer=\frac{(y(y^{-1}-1){)}^n}{n^4}\sum _{Ts}(n^2(y^{-1}-1{)}^{-1}{)}^{|Ts|}\prod _{i=1}^{|Ts|}T{s}_i^{T{s}_i}$$
• 其中 $Ts$ 为一种连通块划分， $|Ts|$ 为连通块个数， $T{s}_i$ 为第 $i$ 个连通块大小， ${\prod }_{i=1}^{|Ts|}T{s}_i^{T{s}_i}$ 即为对划分方案 $Ts$ 来说，每个连通块的生成树个数之积，再乘上每个连通块大小之积的平方
• 就 $\sum$ 后面的东西进行考虑，如果只有一个大小为 $i$ 的连通块，设
• $$Y=n^2(y^{-1}-1{)}^{-1}$$
• $$f[i]=i^i\times Y$$
• 设 $F(x)$ 为 $f[i]$ 的指数生成函数，那么发现 ${\sum }_{Ts}(n^2(y^{-1}-1{)}^{-1}{)}^{|Ts|}{\prod }_{i=1}^{|Ts|}T{s}_i^{T{s}_i}$ 其实就等于 $n![x^n]e^{F(x)}$ ，直接套用多项式 exp 的板子即可
• $O(n\log n)$

Code

#include <bits/stdc++.h>

// 20030830

inline int read()
{
	int res = 0; bool bo = 0; char c;
	while (((c = getchar()) < '0' || c > '9') && c != '-');
	if (c == '-') bo = 1; else res = c - 48;
	while ((c = getchar()) >= '0' && c <= '9')
		res = (res << 3) + (res << 1) + (c - 48);
	return bo ? ~res + 1 : res;
}

template <class T>
inline void Swap(T &a, T &b) {T t = a; a = b; b = t;}

const int N = 1e5 + 5, M = N << 1, L = N * 7, ZZQ = 998244353,
INV2 = 499122177, INV4 = 748683265;

int n, y, op;

int qpow(int a, int b)
{
	int res = 1;
	while (b)
	{
		if (b & 1) res = 1ll * res * a % ZZQ;
		a = 1ll * a * a % ZZQ;
		b >>= 1;
	}
	return res;
}

namespace solve0
{
	int fa[N], ans;
	std::map<int, int> app[N];
	
	int cx(int x)
	{
		if (fa[x] != x) fa[x] = cx(fa[x]);
		return fa[x];
	}
	
	void zm(int x, int y)
	{
		int ix = cx(x), iy = cx(y);
		if (ix != iy) fa[iy] = ix;
	}
	
	void work()
	{
		int u, v;
		for (int i = 1; i <= n; i++) fa[i] = i;
		for (int i = 1; i < n; i++)
			u = read(), v = read(), app[u][v] = app[v][u] = 1;
		for (int i = 1; i < n; i++)
		{
			u = read(); v = read();
			if (app[u][v]) zm(u, v);
		}
		for (int i = 1; i <= n; i++) if (cx(i) == i) ans++;
		printf("%d\n", qpow(y, ans));
	}
};

namespace solve1
{
	int ecnt, nxt[M], adj[N], go[M], f[N], g[N];

	void add_edge(int u, int v)
	{
		nxt[++ecnt] = adj[u]; adj[u] = ecnt; go[ecnt] = v;
		nxt[++ecnt] = adj[v]; adj[v] = ecnt; go[ecnt] = u;
	}
	
	void dfs(int u, int fu)
	{
		f[u] = y; g[u] = 1;
		for (int e = adj[u], v; e; e = nxt[e])
		{
			if ((v = go[e]) == fu) continue;
			dfs(v, u);
			int tf = (1ll * f[u] * f[v] + 1ll * g[u] * f[v]
				+ 1ll * f[u] * g[v]) % ZZQ,
				tg = (1ll * g[u] * f[v] + 1ll * g[u] * g[v]) % ZZQ;
			f[u] = tf; g[u] = tg;
		}
	}
	
	void work()
	{
		int u, v;
		for (int i = 1; i < n; i++)
			u = read(), v = read(), add_edge(u, v);
		if (y == 1) return (void) printf("%d\n", qpow(n, n - 2));
		int st = 1ll * y * (qpow(y, ZZQ - 2) - 1) % ZZQ;
		y = 1ll * qpow(qpow(y, ZZQ - 2) - 1, ZZQ - 2) * n % ZZQ;
		dfs(1, 0);
		printf("%d\n", 1ll * qpow(st, n) * f[1] % ZZQ
			* qpow(1ll * n * n % ZZQ, ZZQ - 2) % ZZQ);
	}
};

namespace solve2
{
	int fac[L], inv[L], invf[L], f[L], yp[L], rev[L], ta[L], tb[L],
	tmpf[L], tmpg[L], da[L], inva[L], tf[L], tg[L], tln[L], ans[L];
	
	void FFT(int n, int *a, int op)
	{
		for (int i = 0; i < n; i++) if (i < rev[i]) Swap(a[i], a[rev[i]]);
		yp[n] = qpow(3, (ZZQ - 1) / n * ((n + op) % n));
		for (int i = n >> 1; i; i >>= 1)
			yp[i] = 1ll * yp[i << 1] * yp[i << 1] % ZZQ;
		for (int k = 1; k < n; k <<= 1)
		{
			int x = yp[k << 1];
			for (int i = 0; i < n; i += k << 1)
			{
				int w = 1;
				for (int j = 0; j < k; j++)
				{
					int u = a[i + j], v = 1ll * w * a[i + j + k] % ZZQ;
					a[i + j] = (u + v) % ZZQ;
					a[i + j + k] = (u - v + ZZQ) % ZZQ;
					w = 1ll * w * x % ZZQ;
				}
			}
		}
	}
	
	void polymul(int n, int *a, int *b, int *res)
	{
		int ff = 1, tot = 0;
		while (ff < (n << 1) - 1) ff <<= 1, tot++;
		for (int i = 0; i < ff; i++)
			rev[i] = (rev[i >> 1] >> 1) | ((i & 1) << tot - 1);
		for (int i = 0; i < ff; i++)
		{
			ta[i] = i < n ? a[i] : 0;
			tb[i] = i < n ? b[i] : 0;
		}
		FFT(ff, ta, 1); FFT(ff, tb, 1);
		for (int i = 0; i < ff; i++) ta[i] = 1ll * ta[i] * tb[i] % ZZQ;
		FFT(ff, ta, -1);
		int gg = qpow(ff, ZZQ - 2);
		for (int i = 0; i < ff; i++) res[i] = 1ll * ta[i] * gg % ZZQ;
	}
	
	void polyinv(int n, int *a, int *res)
	{
		res[0] = qpow(a[0], ZZQ - 2);
		int ff = 4, tot = 2, gg = INV4;
		for (int k = 1; k < n; k <<= 1)
		{
			for (int i = 0; i < ff; i++)
				rev[i] = (rev[i >> 1] >> 1) | ((i & 1) << tot - 1);
			for (int i = 0; i < ff; i++)
			{
				tmpf[i] = i < n && i < (k << 1) ? a[i] : 0;
				tmpg[i] = i < k ? res[i] : 0;
			}
			FFT(ff, tmpf, 1); FFT(ff, tmpg, 1);
			for (int i = 0; i < ff; i++)
				tmpf[i] = 1ll * tmpg[i] * (ZZQ + 2 -
				1ll * tmpf[i] * tmpg[i] % ZZQ) % ZZQ;
			FFT(ff, tmpf, -1);
			for (int i = 0; i < (k << 1); i++)
				res[i] = 1ll * tmpf[i] * gg % ZZQ;
			ff <<= 1; tot++;
			gg = 1ll * gg * INV2 % ZZQ;
		}
	}
	
	void polyln(int n, int *a, int *res)
	{
		for (int i = 0; i < n; i++)
			da[(i + n - 1) % n] = 1ll * i * a[i] % ZZQ;
		polyinv(n, a, inva);
		polymul(n, da, inva, res);
		for (int i = n - 2; i >= 0; i--)
			res[i + 1] = 1ll * inv[i + 1] * res[i] % ZZQ;
		res[0] = 0;
	}
	
	void polyexp(int n, int *a, int *res)
	{
		res[0] = 1;
		int ff = 4, tot = 2, gg = INV4;
		for (int k = 1; k < n; k <<= 1)
		{
			for (int i = 0; i < ff; i++)
				tg[i] = i < k ? res[i] : 0;
			polyln(ff, tg, tln);
			for (int i = 0; i < ff; i++)
				tf[i] = i < n && i < (k << 1) ? a[i] : 0;
			for (int i = 0; i < ff; i++)
				rev[i] = (rev[i >> 1] >> 1) | ((i & 1) << tot - 1);
			for (int i = (k << 1); i < ff; i++) tln[i] = 0;
			FFT(ff, tg, 1); FFT(ff, tf, 1); FFT(ff, tln, 1);
			for (int i = 0; i < ff; i++)
				tf[i] = (1ll - tln[i] + tf[i] + ZZQ) * tg[i] % ZZQ;
			FFT(ff, tf, -1);
			for (int i = 0; i < (k << 1); i++)
				res[i] = 1ll * tf[i] * gg % ZZQ;
			ff <<= 1; tot++;
			gg = 1ll * gg * INV2 % ZZQ;
		}
	}
	
	void work()
	{
		if (y == 1) return (void) printf("%d\n", qpow(n, n - 2 << 1));
		int st = 1ll * y * (qpow(y, ZZQ - 2) - 1) % ZZQ;
		y = 1ll * qpow(qpow(y, ZZQ - 2) - 1, ZZQ - 2) * n % ZZQ * n % ZZQ;
		fac[0] = inv[1] = invf[0] = 1;
		for (int i = 1; i < L; i++)
			fac[i] = 1ll * fac[i - 1] * i % ZZQ;
		for (int i = 2; i < L; i++)
			inv[i] = 1ll * (ZZQ - ZZQ / i) * inv[ZZQ % i] % ZZQ;
		for (int i = 1; i < L; i++)
			invf[i] = 1ll * invf[i - 1] * inv[i] % ZZQ;
		for (int i = 1; i <= n; i++)
			f[i] = 1ll * qpow(i, i) * y % ZZQ * invf[i] % ZZQ;
		polyexp(n + 1, f, ans);
		printf("%d\n", 1ll * qpow(st, n) * ans[n] % ZZQ * fac[n] % ZZQ
			* qpow(n, ZZQ - 5) % ZZQ);
	}
};

int main()
{
	n = read(); y = read(); op = read();
	if (op == 0) solve0::work();
	else if (op == 1) solve1::work();
	else solve2::work();
	return 0;
}