【WC2019】数树【子集反演】【结论】【树形dp】【生成函数】【函数求导】【多项式全家桶】

题意：有两棵基于同一点集的树，点集大小为 $n$ ，两棵树中有 $op$ 棵未确定，可以取所有 $n^{n-2}$ 种可能。给每个点染上 $[1,y]$ 中的一个颜色，要求若 $u$ 到 $v$ 在两棵树上的路径完全相同，那么 $u,v$ 必须同色。求所有方案数之和 模 $998244353$。

$n\le 10^5$

设两棵树的边集为 $E_1,E_2$，那么显然答案就是 $y^{n-|E_1\cap E_2|}$。

我这个都没看出来你信吗

op=0

直接用个 set 模拟即可。

op=1

相当于是给定 $E_1$，对所有 $E_2$ 求答案。

即

$$\sum _{E_2\in \mathrm{Tree}}{y}^{n-|E_1\cap E_2|}$$

其中 $\mathrm{Tree}$ 表示所有构成树的边集的集合。

枚举交集

$$\sum _{S\subseteq E_1}{y}^{n-|S|}\sum _{E_2\in \mathrm{Tree}}[E_1\cap E_2=S]$$

令

$$F(S)=\sum _{E_2\in \mathrm{Tree}}[E_1\cap E_2=S]$$

$$G(S)=[S\subseteq E_1]\sum _{E_2\in \mathrm{Tree}}[S\subseteq E_2]$$

显然有

$$G(S)=\sum _{S\subseteq T}F(T)$$

由基本的容斥原理，有

$$F(S)=\sum _{S\subseteq T}(-1{)}^{|T|-|S|}G(T)$$

我们设一个

$$C(S)=\sum _{E_2\in \mathrm{Tree}}[S\subseteq E_2]$$

即有多少棵树包含了边集 $S$

那么原式可以写成

$$\sum _{S\subseteq E_1}{y}^{n-|S|}\sum _{S\subseteq T\subseteq E_1}(-1{)}^{|T|-|S|}C(T)$$

枚举 $T$

$$\sum _{T\subseteq E_1}C(T)\sum _{S\subseteq T}{y}^{n-|S|}(-1{)}^{|T|-|S|}$$

接下来这步有点迷惑，不过也有其他等价的推法

$$\sum _{T\subseteq E_1}C(T)y^{n-|T|}\sum _{S\subseteq T}(-y{)}^{|T|-|S|}$$

用二项式定理就可以把 $S$ 消掉

$$\sum _{T\subseteq E_1}C(T)y^{n-|T|}(1-y{)}^{|T|}$$

看起来已经很简洁了，现在多半还要上个 dp。

但是这个跨连通块连边的方案很难搞，我们还需要一个这方面的计数工具。我们先手玩一下这个 $C$。

有一个结论：一个有 $k$ 棵树的森林，点集大小为 $n$，其中每棵树的大小为 $a_1,a_2,\dots ,a_k$，则将其连成一棵树的方案数是 $n^{k-2}{\prod }_{i=1}^k{a}_i$。

证明可以手玩矩阵树或者用 prufer 序列，我不太会所以略去了。

那么我们把这个 $C(T)$ 暴力拆开。令 $k=n-|T|$，即断开 $T$ 中的点后的连通块个数。

∑ T ⊆ E ∑ { a } : T n k − 2 ∏ i = 1 k a i y k ( 1 − y ) n − k \sum_{T\subseteq E}\sum_{\{a\}:T}n^{k-2}\prod_{i=1}^ka_iy^k(1-y)^{n-k} T⊆E∑​{a}:T∑​nk−2i=1∏k​ai​yk(1−y)n−k

其中 { a } : T \{a\}:T {a}:T 表示枚举一棵包含 $T$ 的树断掉 $T$ 中的边后形成的 $k$ 个连通块大小为 $a_i$。看不懂也没关系，就是枚举所有可能的连通块。

然后就可以 dp 了。

设 $dp(u,k,s)$ 表示 $u$ 为根的子树中分出了 $k$ 个连通块，其中根所在的连通块大小为 $s$ 的方案数。

这个 $O(n^3)$ 的方程很难继续优化了。但我们注意到我们并不需要对所有 $k$ 求出方案，我们要求的是和它有关的一个指数形式的式子。

通过尝试，可以用01分数规划的思想把它搞成这样：

( 1 − y ) n n 2 ∑ T ⊆ E ∑ { a } : T ∏ i = 1 k n y 1 − y a i \frac{(1-y)^n}{n^2}\sum_{T\subseteq E}\sum_{\{a\}:T}\prod_{i=1}^k\frac{ny}{1-y}a_i n2(1−y)n​T⊆E∑​{a}:T∑​i=1∏k​1−yny​ai​

也就是说，每多出一个连通块，额外产生的贡献就是 $z=\frac{ny}{1-y}$。乘上每个连通块的大小再乘个常数就得到了答案。我们可以边 dp 边计算这个的贡献。

设 $dp(u,k)$ 表示 $u$ 为根的子树中，根所在连通块大小为 $k$，已经产生的总贡献。注意因为根还没有确定，所以根的贡献不计入。

可以写出方程：

$$dp(u,k)\leftarrow dp(u,k)\sum _{i=0}^{\infty }zi\cdot dp(v,i)+\sum _{i=0}^{k}dp(u,i)dp(v,k-i)$$

用树形 dp 的套路可以做到 $O(n^2)$，可是还不够。但注意到这东西又乘常数又卷积的，考虑生成函数。即设

$$d{p}_u(x)=\sum _{i=0}^{\infty }dp(u,i)x^i$$

对于这个 $\sum i\cdot dp(v,i)$，发现它就是 $d{p}_u^{\prime }(1)$（？）。我们设 $f_u=z\cdot d{p}_u^{\prime }(1)$，发现答案就是 $f_1$ 乘上前面那坨常数。

方程可以改写成

$$d{p}_u(x)\leftarrow d{p}_u(x)f_v+d{p}_u(x)d{p}_v(x)$$

为了推出 $f_u$，求个导。

$$d{p}_u^{\prime }(x)\leftarrow d{p}_u^{\prime }(x)f_v+d{p}_u^{\prime }(x)d{p}_v(x)+d{p}_u(x)d{p}_v^{\prime }(x)$$

代入 $x=1$,再乘个 $z$。

$$z\cdot d{p}_u^{\prime }(1)\leftarrow z\cdot d{p}_u^{\prime }(1)f_v+z\cdot d{p}_u^{\prime }(1)d{p}_v(1)+d{p}_u(1)\cdot z\cdot d{p}_v^{\prime }(1)$$

设 $g_u=d{p}_u(1)$，就得到了优美的方程

$$f_u\leftarrow f_u{f}_v+f_u{g}_v+f_v{g}_u$$

$g_u$ 用最开始的方程算

$$g_u\leftarrow g_u(f_v+g_v)$$

就可以 $O(n)$ 求了。

另一种等价的思路是设 $f(u,0/1)$ 表示根结点所在的连通块是否计入了贡献，不过比求导还难想……

op=2

其实比前面还简单点……

用同样的思路推式子

$$\sum _{E_1\in \mathrm{Tree}}\sum _{E_2\in \mathrm{Tree}}{y}^{n-|E_1\cap E_2|}$$

$$=\sum _S{y}^{n-|S|}\sum _{E_1\in \mathrm{Tree}}\sum _{E_2\in \mathrm{Tree}}[E_1\cap E_2=S]$$

$$=\sum _S{y}^{n-|S|}\sum _{S\subseteq T}{C}^2(T)(-1{)}^{|T|-|S|}$$

$$=\sum _T{C}^2(T)\sum _{S\subseteq T}(-1{)}^{|T|-|S|}{y}^{n-|S|}$$

$$=\sum _T{C}^2(T)y^{n-|T|}\sum _{S\subseteq T}(-y{)}^{|T|-|S|}$$

$$=\sum _T{C}^2(T)y^{n-|T|}(1-y{)}^{|T|}$$

= ∑ T ∑ { a } : T n 2 ( k − 2 ) ∏ i = 1 k a i 2 y k ( 1 − y ) n − k =\sum_{T}\sum_{\{a\}:T}n^{2(k-2)}\prod_{i=1}^ka^2_iy^k(1-y)^{n-k} =T∑​{a}:T∑​n2(k−2)i=1∏k​ai2​yk(1−y)n−k

= ( 1 − y ) n n 4 ∑ T ∑ { a } : T ∏ i = 1 k n 2 y 1 − y a i 2 =\frac{(1-y)^n}{n^4}\sum_{T}\sum_{\{a\}:T}\prod_{i=1}^k\frac{n^2y}{1-y}a_i^2 =n4(1−y)n​T∑​{a}:T∑​i=1∏k​1−yn2y​ai2​

考虑一个大小为 $s$ 的连通块，其贡献为后面那坨乘上自己内部的连边方案 $s^{s-2}$，即

$$\frac{n^{2}y}{1-y}{s}^s$$

构造出 EGF 然后 $\exp$ 即可。

#include <iostream>
#include <cstdio>
#include <cstring>
#include <cctype>
#include <utility>
#include <vector>
#include <set>
#define MAXN 262144+5
using namespace std;
inline int read()
{
	int ans=0;
	char c=getchar();
	while (!isdigit(c)) c=getchar();
	while (isdigit(c)) ans=(ans<<3)+(ans<<1)+(c^48),c=getchar();
	return ans;
}
const int MOD=998244353;
inline int add(const int& x,const int& y){return x+y>=MOD? x+y-MOD:x+y;}
inline int dec(const int& x,const int& y){return x<y? x-y+MOD:x-y;}
typedef long long ll;
inline int qpow(int a,int p)
{
	int ans=1;
	while (p)
	{
		if (p&1) ans=(ll)ans*a%MOD;
		a=(ll)a*a%MOD,p>>=1;
	}
	return ans;
}
int n,y;
#define inv(x) qpow(x,MOD-2)
namespace solve1
{
	typedef pair<int,int> pi;
	set<pi> s;
	int main()
	{
		if (y==1) return puts("1"),0;
		for (int i=1;i<n;i++)
		{
			int x,y;
			x=read(),y=read();
			if (x>y) swap(x,y);
			s.insert(make_pair(x,y));
		}
		int t=n;
		for (int i=1;i<n;i++)
		{
			int x,y;
			x=read(),y=read();
			if (x>y) swap(x,y);
			if (s.count(make_pair(x,y))) --t;
		}
		printf("%d\n",qpow(y,t));
		return 0;
	}
}
namespace solve2
{
	vector<int> e[MAXN];
	int F[MAXN],G[MAXN],z;
	void dfs(int u,int f)
	{
		F[u]=z,G[u]=1;
		for (int i=0;i<(int)e[u].size();i++)
		{
			int v=e[u][i];
			if (v!=f)
			{
				dfs(v,u);
				F[u]=((ll)F[u]*F[v]%MOD+(ll)F[u]*G[v]%MOD+(ll)F[v]*G[u]%MOD)%MOD;
				G[u]=(ll)G[u]*add(F[v],G[v])%MOD;
			}			
		}
	}
	int main()
	{
		if (y==1) return printf("%d\n",qpow(n,n-2)),0;
		for (int i=1;i<n;i++)
		{
			int u,v;
			u=read(),v=read();
			e[u].push_back(v),e[v].push_back(u);
		}
		z=(ll)n*y%MOD*inv(dec(1,y))%MOD;
		dfs(1,0);
		printf("%lld\n",(ll)F[1]*qpow(dec(1,y),n)%MOD*inv((ll)n*n%MOD)%MOD);
		return 0;
	}	
}
namespace solve3
{
	int rt[2][24];
	int l,lim,r[MAXN];
	inline void init(){lim=1<<l;for (int i=0;i<lim;i++) r[i]=(r[i>>1]>>1)|((i&1)<<(l-1));}
	inline void NTT(int* a,int type)
	{
		for (int i=0;i<lim;i++) if (i<r[i]) swap(a[i],a[r[i]]);
		for (int L=0;L<l;L++)
		{
			int mid=1<<L,len=mid<<1;
			ll Wn=rt[type][L+1];
			for (int s=0;s<lim;s+=len)
				for (int k=0,w=1;k<mid;k++,w=w*Wn%MOD)
				{
					int x=a[s+k],y=(ll)w*a[s+mid+k]%MOD;
					a[s+k]=add(x,y),a[s+mid+k]=dec(x,y);
				}
		}
		if (type)
		{
			int t=inv(lim);
			for (int i=0;i<lim;i++) a[i]=(ll)a[i]*t%MOD;
		}
	}
	void getinv(int* A,int* B,int n)
	{
		if (n==1) return (void)(*B=1);
		static int f[MAXN],t[MAXN];
		getinv(A,t,(n+1)>>1);
		for (l=0;(1<<l)<=(n<<1);++l);
		init();
		for (int i=0;i<n;i++) f[i]=A[i];
		for (int i=n;i<lim;i++) f[i]=t[i]=0;
		NTT(f,0),NTT(t,0);
		for (int i=0;i<lim;i++) B[i]=(ll)t[i]*dec(2,(ll)f[i]*t[i]%MOD)%MOD;
		NTT(B,1);
		for (int i=n;i<lim;i++) B[i]=0;
	}
	inline void deriv(int* A,int* B,int n){for (int i=0;i<n-1;i++) B[i]=A[i+1]*(i+1ll)%MOD;B[n-1]=0;}
	inline void integ(int* A,int* B,int n){for (int i=1;i<n;i++) B[i]=(ll)A[i-1]*inv(i)%MOD;B[0]=0;}
	inline void getln(int* A,int* B,int n)
	{
		static int f[MAXN],g[MAXN];
		deriv(A,f,n),getinv(A,g,n);
		NTT(f,0),NTT(g,0);
		for (int i=0;i<lim;i++) f[i]=(ll)f[i]*g[i]%MOD;
		NTT(f,1);
		integ(f,B,n);
		for (int i=n;i<lim;i++) B[i]=0;
	}
	void getexp(int* A,int* B,int n)
	{
		if (n==1) return (void)(*B=1);
		static int f[MAXN],g[MAXN];
		getexp(A,f,(n+1)>>1);
		getln(f,g,n);
		for (int i=0;i<n;i++) g[i]=dec(A[i],g[i]);
		++g[0];
		for (int i=n;i<lim;i++) f[i]=g[i]=0;
		NTT(f,0),NTT(g,0);
		for (int i=0;i<lim;i++) B[i]=(ll)f[i]*g[i]%MOD;
		NTT(B,1);
		for (int i=n;i<lim;i++) B[i]=0;
	}
	int F[MAXN],G[MAXN],fac[MAXN],finv[MAXN];
	int main()
	{
		if (y==1) return printf("%d\n",qpow(n,2*n-4)),0;
		rt[0][23]=qpow(3,119),rt[1][23]=inv(rt[0][23]);
		for (int i=22;i>=0;i--)
		{
			rt[0][i]=(ll)rt[0][i+1]*rt[0][i+1]%MOD;
			rt[1][i]=(ll)rt[1][i+1]*rt[1][i+1]%MOD;
		}
		fac[0]=1;
		for (int i=1;i<=n;i++) fac[i]=(ll)fac[i-1]*i%MOD;
		finv[n]=inv(fac[n]);
		for (int i=n-1;i>=0;i--) finv[i]=(ll)finv[i+1]*(i+1)%MOD;
		int t=(ll)n*n%MOD*y%MOD*inv(dec(1,y))%MOD;
		for (int i=1;i<=n;i++) F[i]=(ll)t*qpow(i,i)%MOD*finv[i]%MOD;
		getexp(F,G,n+1);
		printf("%lld\n",(ll)G[n]*qpow(dec(1,y),n)%MOD*qpow(n,MOD-5)%MOD*fac[n]%MOD);
		return 0;
	}	
}
int main()
{
	n=read(),y=read();
	switch(read())
	{
		case 0:return solve1::main();
		case 1:return solve2::main();
		case 2:return solve3::main();
	}
	return 0;
}