洛谷 P5206: bzoj 5475: LOJ 2983: [WC2019] 數樹

一道技巧性非常強的計數題。

題目傳送門：洛谷P5206。

題意簡述：

給定 \(n, y\)。

一張圖有 \(|V| = n\) 個點。對於兩棵樹 \(T_1=G(V, E_1)\) 和 \(T_2=G(V, E_2)\)，定義這兩棵樹的權值 \(F(E_1, E_2)\) 為 \(y\) 的 \(G'=(V,E_1\cap E_2)\) 的連通塊個數次方。

即 \(F(E_1, E_2) = y^{n - |E_1\cap E_2|}\)（因為兩棵樹的邊集的交必然是一個森林，而森林的連通塊個數等於 \(|V| - |E|\)）。

有 \(3\) 類問題：

1. \(\mathrm{op}=0\)，給定 \(E_1, E_2\)，計算 \(F(E_1, E_2)\)。
2. \(\mathrm{op}=1\)，給定 \(E_1\)，計算 \(\displaystyle \sum_{E_2} F(E_1, E_2)\)。
3. \(\mathrm{op}=2\)，計算 \(\displaystyle \sum_{E_1} \sum_{E_2} F(E_1, E_2)\)。

其中 \(\displaystyle \sum_{E}\) 的意義為枚舉所有 \(n^{n - 2}\) 種不同形態的有標號無根樹。

其中 \(n\le 10^5\)。答案對 \(998,244,353\) 取模。

題解：

在這里特別感謝 @rqy 的題解 [WC2019] 數樹。
我對此題的理解以及代碼實現的細節很大程度上借鑒了他的題解。

對於 \(y = 1\)：

\(y = 1\) 時在后續討論中有一些無意義的地方，這里先判掉。
\(y = 1\) 時 \(F(E_1, E_2) = 1\)。

\(\mathrm{op}=0\) 時，只有 \(1\) 種方案，則答案為 \(1\)。
\(\mathrm{op}=1\) 時，有 \(n^{n-2}\) 種方案，則答案為 \(n^{n-2}\)。
\(\mathrm{op}=2\) 時，有 \(n^{2(n-2)}\) 種方案，則答案為 \(n^{2(n-2)}\)。

namespace Solver0 {
	inline LL solve() {
		if (op == 0) return 1;
		if (op == 1) return qPow(N, N - 2);
		if (op == 2) return qPow(N, 2 * (N - 2));
		return 0;
	}
}

對於 \(\mathrm{op} = 0\)：

直接統計出有多少條邊重合，隨便用一個 std::set<std::pair<int,int> > 統計一下。

namespace Solver1 {
	typedef std::pair<int, int> pii;
	std::set<pii> S;
	
	inline LL solve() {
		for (int i = 1, x, y; i < N; ++i) {
			scanf("%d%d", &x, &y);
			if (x > y) std::swap(x, y);
			S.insert(std::make_pair(x, y));
		}
		int cnt = N;
		for (int i = 1, x, y; i < N; ++i) {
			scanf("%d%d", &x, &y);
			if (x > y) std::swap(x, y);
			cnt -= S.count(std::make_pair(x, y));
		}
		return qPow(Y, cnt);
	}
}

對於 \(y\ne 1\) 且 \(\mathrm{op} = 1\)：

給定 \(E_1\)，對於每一種 \(E_2\)，統計 \(y^{n-|E_1\cap E_2|}\)。

考慮枚舉 \(S=E_1\cap E_2\)，則答案為：

\[\begin{aligned}&\sum_{E_2}F(E_1,E_2)\\=&\sum_{E_2}y^{n-|E_1\cap E_2|}\\=&\sum_{S}y^{n-|S|}\sum_{S=E_1\cap E_2}\end{aligned} \]

后面這個不是很優美，我們容斥一下，希望變成 \(S\subseteq E_1\cap E_2\)，即 \(S\subseteq E_1\) 且 \(S\subseteq E_2\) 的形式：

考慮一個容斥原理的式子：\(\displaystyle f(S)=\sum_{T\subseteq S}\sum_{P\subseteq T}(-1)^{|T|-|P|}f(P)\)。

令 \(S=E_1\cap E_2\)，\(f(S)=y^{n-|S|}\)，帶進去一下，得到：

\[\begin{aligned}&\sum_{E_2}y^{n-|E_1\cap E_2|}\\=&\sum_{E_2}\sum_{T\subseteq E_1\cap E_2}\sum_{P\subseteq T}(-1)^{|T|-|P|}y^{n-|P|}\\=&\sum_{T\subseteq E_1}g(T)\sum_{P\subseteq T}(-1)^{|T|-|P|}y^{n-|P|}\\=&\sum_{T\subseteq E_1}g(T)y^{n-|T|}\sum_{P\subseteq T}(-y)^{|T|-|P|}\\=&\sum_{T\subseteq E_1}g(T)y^{n-|T|}\sum_{k=0}^{|T|}\binom{|T|}{k}1^{k}(-y)^{|T|-k}\\=&\sum_{T\subseteq E_1}g(T)y^{n-|T|}(1-y)^{|T|}\end{aligned} \]

其中 \(g(T)\) 表示包含邊集 \(T\) 的樹個數。

這里有一個定理，可以用 \(\mathrm{Pr\ddot{u}fer}\) 序列證明，也可以用 \(\mathrm{Matrix-Tree}\) 定理證明：

• 對於 \(n\) 個點的森林，假設有 \(k\) 個連通分量，每個連通分量的大小分別是 \(a_i\)，則包含這個森林的樹個數為 \(\displaystyle n^{k-2}\prod_{i=1}^{k}a_i\)。

再帶入可得：

\[\begin{aligned} &\sum_{T\subseteq E_1}g(T)y^{n-|T|}(1-y)^{|T|}\\=&\sum_{T\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_{T\subseteq E_1}\prod_{i=1}^{k}\frac{ny}{1-y}a_i\end{aligned} \]

即一個大小為 \(a\) 的連通分量為答案貢獻一個乘積 \(Ka\)，其中 \(\displaystyle K = \frac{ny}{1-y}\)。

據此我們可以寫出一個 \(\mathrm{DP}\) 式子，用 \(\mathrm{f}[u][i]\) 表示 \(u\) 的子樹中，\(u\) 所在的連通分量大小為 \(i\) 的答案。

\(i\) 所在的連通分量還未完整，所以不計入答案，而是存入狀態，轉類似於樹上背包，轉移比較顯然，這種做法是 \(\mathrm{O}(n^2)\) 的。

我們可以考慮每個貢獻的組合意義，大小為 \(a_i\) 的連通分量貢獻相當於在這個連通分量中選取一個點產生 \(K\) 的乘積貢獻。

據此我們可以優化狀態表示：\(\mathrm{f}[u][0/1]\) 表示 \(u\) 的子樹中，當前連通分量是否已經做出貢獻，這里為了方便轉移，當前連通分量如果已經做出貢獻就計入答案。

具體轉移可以參考下面的代碼。

namespace Solver2 {
	int h[MN], nxt[MN * 2], to[MN * 2], tot;
	inline void ins(int x, int y) {
		nxt[++tot] = h[x], to[tot] = y, h[x] = tot;
	}
	
	LL K, f[MN], g[MN];
	void DFS(int u, int fz) {
		f[u] = 1, g[u] = K;
		for (int i = h[u]; i; i = nxt[i]) if (to[i] != fz) {
			DFS(to[i], u);
			g[u] = (f[u] * g[to[i]] + g[u] * f[to[i]] + g[u] * g[to[i]]) % Mod;
			f[u] = (f[u] * f[to[i]] + f[u] * g[to[i]]) % Mod;
		}
	}
	
	inline LL solve() {
		for (int i = 1, x, y; i < N; ++i) {
			scanf("%d%d", &x, &y);
			ins(x, y), ins(y, x);
		}
		K = (LL)N * Y % Mod * qPow(1 - Y, Mod - 2) % Mod;
		LL coef = qPow(1 - Y, N) * qPow(N, Mod - 3) % Mod;
		DFS(1, 0);
		return g[1] * coef % Mod;
	}
}

對於 \(y\ne 1\) 且 \(\mathrm{op} = 2\)：

類似 \(\mathrm{op} = 1\)，我們可以寫出下面的式子：

\[\begin{aligned}&\sum_{E_1}\sum_{E_2}F(E_1,E_2)\\=&\sum_{E_1}\sum_{E_2}y^{n-|E_1\cap E_2|}\\=&\sum_{E_1}\sum_{E_2}\sum_{T\subseteq E_1\cap E_2}\sum_{P\subseteq T}(-1)^{|T|-|P|}y^{n-|P|}\\=&\sum_{T}g^2(T)\sum_{P\subseteq T}(-1)^{|T|-|P|}y^{n-|P|}\\=&\sum_{T}g^2(T)y^{n-|T|}\sum_{P\subseteq T}(-1)^{|T|-|P|}y^{|T|-|P|}\\=&\sum_{T}g^2(T)y^{n-|T|}\sum_{P\subseteq T}(-y)^{|T|-|P|}\\=&\sum_{T}g^2(T)y^{n-|T|}\sum_{k=0}^{|T|}\binom{|T|}{k}1^{k}(-y)^{|T|-k}\\=&\sum_{T}g^2(T)y^{n-|T|}(1-y)^{|T|}\\=&\sum_{T}y^{k}(1-y)^{n-k}n^{2(k-2)}\prod_{i=1}^{k}a_i^2\\=&\frac{(1-y)^n}{n^4}\sum_{T}\prod_{i=1}^{k}\frac{n^2y}{1-y}a_i^2\end{aligned} \]

此時我們換個角度考慮，與其考慮邊集 \(T\)，不如考慮每個連通分量。

這里的每個連通分量是無序的，所以相當於 \(n\) 個有標號小球扔進 \(k\) 個無標號盒子，不允許空盒子。
而一個裝有 \(a\) 個球的盒子，每個生成樹都會貢獻 \(\displaystyle\frac{n^2y}{1-y}a^2\) 作為乘積。
而 \(a\) 個點的生成樹有 \(a^{a-2}\) 個，所以總共貢獻 \(\displaystyle\frac{n^2y}{1-y}a^a\)。

對生成函數比較敏感的同學能夠看出，總方案的指數型生成函數就是單個盒子的指數型生成函數的 \(\mathrm{Exp}\)，因為對應了集合與元素的關系。
即 \(\displaystyle\mathrm{f}=\left\{\frac{n^2y}{1-y}i^i\right\}_{i=1}^{\infty}\) 對應的指數型生成函數 \(\displaystyle\mathrm{F}(x)=\sum_{i=1}^{\infty}\frac{n^2y}{(1-y)i!}i^{i}x^{i}\)。
令 \(\mathrm{G}=e^{\mathrm{F}}\)，則指數型生成函數 \(\mathrm{G}\) 對應的序列的第 \(n\) 項便是所求答案，即答案為 \(n!\cdot [x^n]\mathrm{G}\)。

如果你對生成函數比較不敏感，也可以推式子得到：
對於有標號球有標號盒子記錄貢獻，可以使用指數型生成函數的冪次表示：
\(n\) 個有標號球，\(k\) 個有標號盒子就是 \([x^n]\mathrm{F}^k\)，但是這里盒子是無標號的，所以要除以 \(k!\)。
那么最終的答案是 \(\displaystyle\frac{(1-y)^n\cdot n!}{n^4}[x^n]\sum_{k=1}^{n}\frac{1}{k!}\left(\sum_{i=1}^{\infty}\frac{n^2y}{(1-y)\cdot i!}i^ix^i\right)^k\)。
根據 \(e^x\) 的泰勒展開公式，可以發現形式是完全相同的。

那么直接套多項式 \(\mathrm{Exp}\) 模板即可。

#include <cstdio>
#include <cstring>
#include <algorithm>
#include <set>

typedef long long LL;
const int MN = 100005;
const LL Mod = 998244353;

inline LL qPow(LL b, int e) {
	LL a = 1;
	for (; e; b = b * b % Mod, e >>= 1)
		if (e & 1) a = a * b % Mod;
	return a;
}

int N, op;
LL Y;

namespace Solver0 {
	inline LL solve() {
		if (op == 0) return 1;
		if (op == 1) return qPow(N, N - 2);
		if (op == 2) return qPow(N, 2 * (N - 2));
		return 0;
	}
}

namespace Solver1 {
	typedef std::pair<int, int> pii;
	std::set<pii> S;
	
	inline LL solve() {
		for (int i = 1, x, y; i < N; ++i) {
			scanf("%d%d", &x, &y);
			if (x > y) std::swap(x, y);
			S.insert(std::make_pair(x, y));
		}
		int cnt = N;
		for (int i = 1, x, y; i < N; ++i) {
			scanf("%d%d", &x, &y);
			if (x > y) std::swap(x, y);
			cnt -= S.count(std::make_pair(x, y));
		}
		return qPow(Y, cnt);
	}
}

namespace Solver2 {
	int h[MN], nxt[MN * 2], to[MN * 2], tot;
	inline void ins(int x, int y) {
		nxt[++tot] = h[x], to[tot] = y, h[x] = tot;
	}
	
	LL K, f[MN], g[MN];
	void DFS(int u, int fz) {
		f[u] = 1, g[u] = K;
		for (int i = h[u]; i; i = nxt[i]) if (to[i] != fz) {
			DFS(to[i], u);
			g[u] = (f[u] * g[to[i]] + g[u] * f[to[i]] + g[u] * g[to[i]]) % Mod;
			f[u] = (f[u] * f[to[i]] + f[u] * g[to[i]]) % Mod;
		}
	}
	
	inline LL solve() {
		for (int i = 1, x, y; i < N; ++i) {
			scanf("%d%d", &x, &y);
			ins(x, y), ins(y, x);
		}
		K = (LL)N * Y % Mod * qPow(1 - Y, Mod - 2) % Mod;
		LL coef = qPow(1 - Y, N) * qPow(N, Mod - 3) % Mod;
		DFS(1, 0);
		return g[1] * coef % Mod;
	}
}

namespace Solver3 {
	const int MS = 1 << 19;
	const int G = 3, iG = 332748118;
	int Sz = 0, R[MS]; LL InvSz;
	LL Inv[MS], Fac[MS], iFac[MS];
	
	inline void Init() {
		Inv[1] = 1;
		for (int i = 2; i < MS; ++i) Inv[i] = -(Mod / i) * Inv[Mod % i] % Mod;
		Fac[0] = iFac[0] = 1;
		for (int i = 1; i < MS; ++i) {
			Fac[i] = Fac[i - 1] * i % Mod;
			iFac[i] = iFac[i - 1] * Inv[i] % Mod;
		}
	}
	
	inline void InitFNTT(int n) {
		int Bt = 0;
		for (; 1 << Bt < n; ++Bt) ;
		if ((1 << Bt) == Sz) return ;
		Sz = 1 << Bt, InvSz = -(Mod - 1) / Sz;
		for (int i = 1; i < Sz; ++i) R[i] = R[i >> 1] >> 1 | (i & 1) << (Bt - 1);
	}
	
	inline void FNTT(LL *A, int Type) {
		for (int i = 0; i < Sz; ++i) if (R[i] < i) std::swap(A[R[i]], A[i]);
		for (int j = 1, j2 = 2; j < Sz; j <<= 1, j2 <<= 1) {
			LL gn = qPow(~Type ? G : iG, (Mod - 1) / j2), g;
			for (int i = 0, k; i < Sz; i += j2) {
				for (k = 0, g = 1; k < j; ++k, g = g * gn % Mod) {
					LL X = A[i + k], Y = g * A[i + j + k] % Mod;
					A[i + k] = (X + Y) % Mod, A[i + j + k] = (X - Y) % Mod;
				}
			}
		}
		if (Type == -1) for (int i = 0; i < Sz; ++i) A[i] = A[i] * InvSz % Mod;
	}
	
	inline void PolyInv(LL *A, int N, LL *B) {
		B[0] = qPow(A[0], Mod - 2);
		static LL tA[MS], tB[MS];
		for (int L = 1; L < N; L <<= 1) {
			int L2 = L << 1, L4 = L << 2;
			InitFNTT(L4);
			for (int i = 0; i < L2; ++i) tA[i] = A[i];
			for (int i = L2; i < Sz; ++i) tA[i] = 0;
			for (int i = 0; i < L; ++i) tB[i] = B[i];
			for (int i = L; i < Sz; ++i) tB[i] = 0;
			FNTT(tA, 1), FNTT(tB, 1);
			for (int i = 0; i < Sz; ++i) tB[i] = (2 - tA[i] * tB[i]) % Mod * tB[i] % Mod;
			FNTT(tB, -1);
			for (int i = 0; i < L2; ++i) B[i] = tB[i];
		}
	}
	
	inline void PolyLn(LL *A, int N, LL *B) {
		static LL tA[MS], tB[MS];
		PolyInv(A, N, tB);
		InitFNTT(N * 2);
		for (int i = 1; i < N; ++i) tA[i - 1] = A[i] * i % Mod;
		for (int i = N - 1; i < Sz; ++i) tA[i] = 0;
		for (int i = N; i < Sz; ++i) tB[i] = 0;
		FNTT(tA, 1), FNTT(tB, 1);
		for (int i = 0; i < Sz; ++i) tA[i] = tA[i] * tB[i] % Mod;
		FNTT(tA, -1);
		B[0] = 0;
		for (int i = 1; i < N; ++i) B[i] = tA[i - 1] * Inv[i] % Mod;
	}
	
	inline void PolyExp(LL *A, int N, LL *B) {
		B[0] = 1;
		static LL tA[MS], tB[MS];
		for (int L = 1; L < N; L <<= 1) {
			int L2 = L << 1, L4 = L << 2;
			PolyLn(B, L2, tA);
			InitFNTT(L4);
			for (int i = 0; i < L2; ++i) tA[i] = (!i + A[i] - tA[i]) % Mod;
			for (int i = L2; i < Sz; ++i) tA[i] = 0;
			for (int i = 0; i < L2; ++i) tB[i] = B[i];
			for (int i = L2; i < Sz; ++i) tB[i] = 0;
			FNTT(tA, 1), FNTT(tB, 1);
			for (int i = 0; i < Sz; ++i) tA[i] = tA[i] * tB[i] % Mod;
			FNTT(tA, -1);
			for (int i = 0; i < L2; ++i) B[i] = tA[i];
		}
	}
	
	inline LL solve() {
		Init();
		LL K = (LL)N * N % Mod * Y % Mod * qPow(1 - Y, Mod - 2) % Mod;
		LL coef = qPow(1 - Y, N) * qPow(N, Mod - 5) % Mod;
		static LL A[MS], B[MS];
		A[0] = 0;
		for (int i = 1; i <= N; ++i) A[i] = K * qPow(i, i) % Mod * iFac[i] % Mod;
		PolyExp(A, N + 1, B);
		return coef * (B[N] * Fac[N] % Mod) % Mod;
	}
}

int main() {
	scanf("%d%lld%d", &N, &Y, &op);
	if (Y == 1) printf("%lld\n", Solver0::solve());
	else if (op == 0) printf("%lld\n", (Solver1::solve() + Mod) % Mod);
	else if (op == 1) printf("%lld\n", (Solver2::solve() + Mod) % Mod);
	else if (op == 2) printf("%lld\n", (Solver3::solve() + Mod) % Mod);
	return 0;
}