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

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

$n≤105n\leq 10^5$

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

~~我这个都没看出来你信吗~~

op=0

直接用个 set 模拟即可。

op=1

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

即

$∑E2∈Tree⁡yn−∣E1∩E2∣\sum_{E_2\in \operatorname{Tree}} y^{n-|E_1\cap E_2|}$

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

枚举交集

$∑S⊆E1yn−∣S∣∑E2∈Tree⁡[E1∩E2=S]\sum_{S\subseteq E_1}y^{n-|S|}\sum_{E_2\in \operatorname{Tree}}[E_1\cap E_2=S]$

令

$F(S)=∑E2∈Tree⁡[E1∩E2=S]F(S)=\sum_{E_2\in \operatorname{Tree}}[E_1\cap E_2=S]$

$\subseteq E_1]\sum_{E_2\in \operatorname{Tree}}[S\subseteq E_2]$

显然有

$G(S)=∑S⊆TF(T)G(S)=\sum_{S\subseteq T}F(T)$

由基本的容斥原理，有

$F(S)=∑S⊆T(−1)∣T∣−∣S∣G(T)F(S)=\sum_{S\subseteq T}(-1)^{|T|-|S|}G(T)$

我们设一个

$C(S)=∑E2∈Tree⁡[S⊆E2]C(S)=\sum_{E_2\in \operatorname{Tree}}[S\subseteq E_2]$

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

那么原式可以写成

$∑S⊆E1yn−∣S∣∑S⊆T⊆E1(−1)∣T∣−∣S∣C(T)\sum_{S\subseteq E_1}y^{n-|S|}\sum_{S\subseteq T \subseteq E_1}(-1)^{|T|-|S|}C(T)$

枚举 $T$

$∑T⊆E1C(T)∑S⊆Tyn−∣S∣(−1)∣T∣−∣S∣\sum_{T\subseteq E_1}C(T)\sum_{S\subseteq T}y^{n-|S|}(-1)^{|T|-|S|}$

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

$∑T⊆E1C(T)yn−∣T∣∑S⊆T(−y)∣T∣−∣S∣\sum_{T\subseteq E_1}C(T)y^{n-|T|}\sum_{S\subseteq T}(-y)^{|T|-|S|}$

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

$∑T⊆E1C(T)yn−∣T∣(1−y)∣T∣\sum_{T\subseteq E_1}C(T)y^{n-|T|}(1-y)^{|T|}$

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

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

有一个结论：一个有 $k$ 棵树的森林，点集大小为 $n$ ，其中每棵树的大小为 $a1,a2,…,aka_1,a_2,\dots,a_k$ ，则将其连成一棵树的方案数是 $nk−2∏i=1kain^{k-2}\prod_{i=1}^ka_i$ 。

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

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

$∑T⊆E∑{a}:Tnk−2∏i=1kaiyk(1−y)n−k\sum_{T\subseteq E}\sum_{\{a\}:T}n^{k-2}\prod_{i=1}^ka_iy^k(1-y)^{n-k}$

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

然后就可以 dp 了。

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

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

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

$(1−y)nn2∑T⊆E∑{a}:T∏i=1kny1−yai\frac{(1-y)^n}{n^2}\sum_{T\subseteq E}\sum_{\{a\}:T}\prod_{i=1}^k\frac{ny}{1-y}a_i$

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

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

可以写出方程：

$dp(u,k)←dp(u,k)∑i=0∞zi⋅dp(v,i)+∑i=0kdp(u,i)dp(v,k−i)dp(u,k)\leftarrow dp(u,k)\sum_{i=0}^{\infin}zi\cdot dp(v,i)+\sum_{i=0}^kdp(u,i)dp(v,k-i)$

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

$dpu(x)=∑i=0∞dp(u,i)xidp_u(x)=\sum_{i=0}^{\infin}dp(u,i)x^i$

对于这个 $∑i⋅dp(v,i)\sum i\cdot dp(v,i)$ ，发现它就是 $dp_u'(1)$ （？）。我们设 $fu=z⋅dpu′(1)f_u=z\cdot dp_u'(1)$ ，发现答案就是 $f_1$ 乘上前面那坨常数。

方程可以改写成

$dpu(x)←dpu(x)fv+dpu(x)dpv(x)dp_u(x)\leftarrow dp_u(x)f_v+dp_u(x)dp_v(x)$

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

$dpu′(x)←dpu′(x)fv+dpu′(x)dpv(x)+dpu(x)dpv′(x)dp_u'(x)\leftarrow dp_u'(x)f_v+dp_u'(x)dp_v(x)+dp_u(x)dp_v'(x)$

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

$z⋅dpu′(1)←z⋅dpu′(1)fv+z⋅dpu′(1)dpv(1)+dpu(1)⋅z⋅dpv′(1)z\cdot dp_u'(1)\leftarrow z\cdot dp_u'(1)f_v+z\cdot dp_u'(1)dp_v(1)+dp_u(1)\cdot z\cdot dp_v'(1)$

设 $g_u=dp_u(1)$ ，就得到了优美的方程

$fu←fufv+fugv+fvguf_u\leftarrow f_uf_v+f_ug_v+f_vg_u$

$g_u$ 用最开始的方程算

$gu←gu(fv+gv)g_u\leftarrow g_u(f_v+g_v)$

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

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

op=2

其实比前面还简单点……

用同样的思路推式子

$∑E1∈Tree⁡∑E2∈Tree⁡yn−∣E1∩E2∣\sum_{E_1\in \operatorname{Tree}}\sum_{E_2\in \operatorname{Tree}}y^{n-|E_1\cap E_2|}$

$=∑Syn−∣S∣∑E1∈Tree⁡∑E2∈Tree⁡[E1∩E2=S]=\sum_{S}y^{n-|S|}\sum_{E_1\in \operatorname{Tree}}\sum_{E_2\in \operatorname{Tree}}[E_1\cap E_2=S]$

$=∑Syn−∣S∣∑S⊆TC2(T)(−1)∣T∣−∣S∣=\sum_{S}y^{n-|S|}\sum_{S\subseteq T}C^2(T)(-1)^{|T|-|S|}$

$=∑TC2(T)∑S⊆T(−1)∣T∣−∣S∣yn−∣S∣=\sum_{T}C^2(T)\sum_{S\subseteq T}(-1)^{|T|-|S|}y^{n-|S|}$

$=∑TC2(T)yn−∣T∣∑S⊆T(−y)∣T∣−∣S∣=\sum_{T}C^2(T)y^{n-|T|}\sum_{S\subseteq T}(-y)^{|T|-|S|}$

$=∑TC2(T)yn−∣T∣(1−y)∣T∣=\sum_{T}C^2(T)y^{n-|T|}(1-y)^{|T|}$

$=∑T∑{a}:Tn2(k−2)∏i=1kai2yk(1−y)n−k=\sum_{T}\sum_{\{a\}:T}n^{2(k-2)}\prod_{i=1}^ka^2_iy^k(1-y)^{n-k}$

$=(1−y)nn4∑T∑{a}:T∏i=1kn2y1−yai2=\frac{(1-y)^n}{n^4}\sum_{T}\sum_{\{a\}:T}\prod_{i=1}^k\frac{n^2y}{1-y}a_i^2$

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

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

构造出 EGF 然后 $exp⁡\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;
}