【LOJ #2983】「WC2019」数树

Description

题目链接：

• UOJ #457
• LOJ #2983
• Luogu #5206

本题包含三个问题：

• 问题 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)

时空限制：( exttt{4s/512MB})

Solution

注意 (y=1) 的情况会带来一些困扰，提前判掉即可。

问题 0

根据题意，由两棵树边集的交集构成的同一个连通块内的点均要取同一个颜色，森林的连通块数等于点数减边数，因此求出边集交集的大小 (k)，答案就是 (y^{n-k})。

可以用 ( ext{std::map}) 实现这个过程。

问题 1

暴力矩阵树

假设每个生成树的权值为 (y^{n-k})，(k) 表示这棵生成树和给定的树的边集交集大小，那么要求的就是所有生成树的权值之和。不妨将 (y^n) 提取出来，那么每条原树的边的边权都是 (y^{-1})，不是原树的边的边权就是 (1)。

直接套用矩阵树定理即可。时间复杂度 (mathcal O(n^3))。

听说这个矩阵比较有性质可以优化到 (mathcal O(n))，因为我太菜无法理解题解的做法，所以这种做法就到此为止了。

推式子

不妨将第一棵树的边集叫做 (E_1)，第二棵树的边集叫做 (E_2)，那么就有下面这个式子成立

[egin{aligned} mathrm{ans}&=sum_{E_2}y^{n-|E_1cap E_2|} \ &=sum_{Ssubseteq E_1}y^{n-|S|}sum_{E_2}[S=E_1cap E_2] end{aligned} ]

不妨设

[f(S)=sum_{E_2}[S=E_1cap E_2] ]

那么 (f(S)) 显然不太好求，不妨考虑一下 (f(S)) 的实际含义，(f(S)) 就表示将强制 (E_2) 首先要选边集 (S) 的边，并且强制 (E_2) 不能选 (E_1setminus S) 的边。

前者比较好处理，但是后者不太好处理。我们考虑先去掉后面那个条件，就想到设

[g(S)=sum_{E_2}[Ssubseteq E_1cap E_2] ]

显然有

[g(S)=sum_{Ssubseteq T}f(T) ]

容斥一下就是

[f(S)=sum_{Ssubseteq T}(-1)^{|T|-|S|}g(T) ]

所以

[egin{aligned} mathrm{ans}&=sum_{Ssubseteq E_1}y^{n-|S|}f(S)\ &=sum_{Ssubseteq E_1}y^{n-|S|}sum_{Ssubseteq Tsubseteq E_1}(-1)^{|T|-|S|}g(T)\ &=sum_{Tsubseteq E_1}g(T)sum_{Ssubseteq T}y^{n-|S|}(-1)^{|T|-|S|}\ &=y^nsum_{Tsubseteq E_1}g(T)sum_{i=0}^{|T|}inom{|T|}{i}y^{-i}(-1)^{|T|-i}\ &=y^nsum_{Tsubseteq 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_icdot 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})。具体地

[egin{aligned} g(T)&=sum_{{p_{k-2}}}prod_{i=1}^ka_i^{q_i+1}\ &=prod_{i=1}^ka_isum_{{p_{k-2}}}prod_{i=1}^ka_i^{q_i}\ &=prod_{i=1}^ka_isum_{{p_{k-2}}}prod_{i=1}^{k-2}a_{p_i}\ &=prod_{i=1}^ka_iprod_{i=1}^{k-2}sum_{j=1}^ka_{j}\ &=n^{k-2}prod_{i=1}^ka_i end{aligned} ]

矩阵树定理推导

去掉最后一行最后一列的基尔霍夫矩阵的行列式是

[egin{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 egin{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 egin{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 egin{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 egin{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) 后得到的所有连通块的大小之积。

[egin{aligned} mathrm{ans}&=y^nsum_{Tsubseteq E_1}g(T)(y^{-1}-1)^{|T|}\ &=y^nsum_{Tsubseteq E_1}n^{n-|T|-2}R(T)(y^{-1}-1)^{|T|}\ &=frac{y^nn^n}{n^2}sum_{Tsubseteq E_1}R(T)left(frac{y^{-1}-1}{n} ight)^{|T|}\ end{aligned} ]

我们不妨设 (a(T)_i) 表示 (T) 得到的第 (i) 个连通块大小，那么有

[egin{aligned} mathrm{ans}&=frac{y^nn^n}{n^2}sum_{Tsubseteq E_1}left(frac{y^{-1}-1}{n} ight)^{|T|}prod_{i=1}^{n-|T|}a(T)_i\ &=frac{(1-y)^n}{n^2}sum_{Tsubseteq E_1}prod_{i=1}^{n-|T|}a(T)_ileft(frac{n}{y^{-1}-1} ight) end{aligned} ]

树形 DP

不妨设 (K=frac{n}{y^{-1}-1})，现在只需要考虑计算

[sum_{Tsubseteq E_1}prod_{i=1}^{n-|T|}Kcdot 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

推式子是类似的，不过这时候就没有限制 (Tsubseteq E_1) 了，(T) 可以选取任意无环无重边边集。

[mathrm{ans}=y^nsum_{T}g(T)^2(y^{-1}-1)^{|T|} ]

代入 (g(T)) 的表达式我们有

[egin{aligned} mathrm{ans}&=(1-y)^nsum_{T}R(T)^2n^{2(n-|T|-2)}left(frac{1}{y^{-1}-1} ight)^{n-|T|}\ &=frac{(1-y)^n}{n^4}sum_{T}R(T)^2left(frac{n^2}{y^{-1}-1} ight)^{n-|T|}\ &=frac{(1-y)^n}{n^4}sum_{T}prod_{i=1}^{n-|T|}a(T)_i^2left(frac{n^2}{y^{-1}-1} ight)\ end{aligned} ]

与问题 1 不同的地方体现在，问题 2 没有给出一棵树，因此每种连通块划分方案对应着多个边集，需要在每个连通块内部任意连一棵树出来。而且有标号使得划分连通块大小后还需要给每个连通块分配标号。

因此这涉及到有标号组合对象的拼接，我们考虑指数生成函数，设 (F(x)) 表示每个连通块的指数生成函数，那么

[egin{aligned} F(x)&=sum_{i=1}^{+infty}frac{i^2cdot i^{i-2}cdot n^2}{y^{-1}-1} cdotfrac{x^i}{i!}\ &=sum_{i=1}^{+infty}frac{i^icdot n^2}{y^{-1}-1} cdotfrac{x^i}{i!} end{aligned} ]

那么就有

[egin{aligned} mathrm{ans}&=frac{(1-y)^nn!}{n^4}[x^n]sum_{k=1}^nfrac{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) << '
'; 
	}
}

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 << '
'; 
	}
}

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 << '
'; 
	}
}

int main()
{
	read(n), read(y), read(opt); 

	if (opt == 0)
		task0::answer(); 
	else
	{
		if (y == 1)
			std::cout << qpow(n, opt * (n - 2)) << '
'; 
		else
		{
			if (opt == 1)
				task1::answer(); 
			else
				task2::answer(); 
		}
	}
	return 0; 
}