题意：有两棵基于同一点集的树，点集大小为 n n n ，两棵树中有 o p op op 棵未确定，可以取所有 n n − 2 n^{n-2} nn−2 种可能。给每个点染上 [ 1 , y ] [1,y] [1,y] 中的一个颜色，要求若 u u u 到 v v v 在两棵树上的路径完全相同，那么 u , v u,v u,v 必须同色。求所有方案数之和 模 998244353 998244353 998244353。
n ≤ 1 0 5 n\leq 10^5 n≤105
设两棵树的边集为 E 1 , E 2 E_1,E_2 E1,E2，那么显然答案就是 y n − ∣ E 1 ∩ E 2 ∣ y^{n-|E_1\cap E_2|} yn−∣E1∩E2∣。
我这个都没看出来你信吗
op=0
直接用个 set 模拟即可。
op=1
相当于是给定 E 1 E_1 E1，对所有 E 2 E_2 E2 求答案。
即
∑ E 2 ∈ Tree y n − ∣ E 1 ∩ E 2 ∣ \sum_{E_2\in \operatorname{Tree}} y^{n-|E_1\cap E_2|} E2∈Tree∑yn−∣E1∩E2∣
其中 Tree \operatorname{Tree} Tree 表示所有构成树的边集的集合。
枚举交集
∑ S ⊆ E 1 y n − ∣ S ∣ ∑ E 2 ∈ Tree [ E 1 ∩ E 2 = S ] \sum_{S\subseteq E_1}y^{n-|S|}\sum_{E_2\in \operatorname{Tree}}[E_1\cap E_2=S] S⊆E1∑yn−∣S∣E2∈Tree∑[E1∩E2=S]
令
F ( S ) = ∑ E 2 ∈ Tree [ E 1 ∩ E 2 = S ] F(S)=\sum_{E_2\in \operatorname{Tree}}[E_1\cap E_2=S] F(S)=E2∈Tree∑[E1∩E2=S]
G ( S ) = [ S ⊆ E 1 ] ∑ E 2 ∈ Tree [ S ⊆ E 2 ] G(S)=[S \subseteq E_1]\sum_{E_2\in \operatorname{Tree}}[S\subseteq E_2] G(S)=[S⊆E1]E2∈Tree∑[S⊆E2]
显然有
G ( S ) = ∑ S ⊆ T F ( T ) G(S)=\sum_{S\subseteq T}F(T) G(S)=S⊆T∑F(T)
由基本的容斥原理，有
F ( S ) = ∑ S ⊆ T ( − 1 ) ∣ T ∣ − ∣ S ∣ G ( T ) F(S)=\sum_{S\subseteq T}(-1)^{|T|-|S|}G(T) F(S)=S⊆T∑(−1)∣T∣−∣S∣G(T)
我们设一个
C ( S ) = ∑ E 2 ∈ Tree [ S ⊆ E 2 ] C(S)=\sum_{E_2\in \operatorname{Tree}}[S\subseteq E_2] C(S)=E2∈Tree∑[S⊆E2]
即有多少棵树包含了边集 S S S
那么原式可以写成
∑ S ⊆ E 1 y n − ∣ S ∣ ∑ S ⊆ T ⊆ E 1 ( − 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) S⊆E1∑yn−∣S∣S⊆T⊆E1∑(−1)∣T∣−∣S∣C(T)
枚举 T T T
∑ T ⊆ E 1 C ( T ) ∑ S ⊆ T y n − ∣ S ∣ ( − 1 ) ∣ T ∣ − ∣ S ∣ \sum_{T\subseteq E_1}C(T)\sum_{S\subseteq T}y^{n-|S|}(-1)^{|T|-|S|} T⊆E1∑C(T)S⊆T∑yn−∣S∣(−1)∣T∣−∣S∣
接下来这步有点迷惑，不过也有其他等价的推法
∑ T ⊆ E 1 C ( T ) y n − ∣ T ∣ ∑ S ⊆ T ( − y ) ∣ T ∣ − ∣ S ∣ \sum_{T\subseteq E_1}C(T)y^{n-|T|}\sum_{S\subseteq T}(-y)^{|T|-|S|} T⊆E1∑C(T)yn−∣T∣S⊆T∑(−y)∣T∣−∣S∣
用二项式定理就可以把 S S S 消掉
∑ T ⊆ E 1 C ( T ) y n − ∣ T ∣ ( 1 − y ) ∣ T ∣ \sum_{T\subseteq E_1}C(T)y^{n-|T|}(1-y)^{|T|} T⊆E1∑C(T)yn−∣T∣(1−y)∣T∣
看起来已经很简洁了，现在多半还要上个 dp。
但是这个跨连通块连边的方案很难搞，我们还需要一个这方面的计数工具。我们先手玩一下这个 C C C。
有一个结论：一个有 k k k 棵树的森林，点集大小为 n n n，其中每棵树的大小为 a 1 , a 2 , … , a k a_1,a_2,\dots,a_k a1,a2,…,ak，则将其连成一棵树的方案数是 n k − 2 ∏ i = 1 k a i n^{k-2}\prod_{i=1}^ka_i nk−2∏i=1kai。
证明可以手玩矩阵树或者用 prufer 序列，我不太会所以略去了。
那么我们把这个 C ( T ) C(T) C(T) 暴力拆开。令 k = n − ∣ T ∣ k=n-|T| k=n−∣T∣，即断开 T 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∏kaiyk(1−y)n−k
其中 { a } : T \{a\}:T {a}:T 表示枚举一棵包含 T T T 的树断掉 T T T 中的边后形成的 k k k 个连通块大小为 a i a_i ai。看不懂也没关系，就是枚举所有可能的连通块。
然后就可以 dp 了。
设 d p ( u , k , s ) dp(u,k,s) dp(u,k,s) 表示 u u u 为根的子树中分出了 k k k 个连通块，其中根所在的连通块大小为 s s s 的方案数。
这个 O ( n 3 ) O(n^3) O(n3) 的方程很难继续优化了。但我们注意到我们并不需要对所有 k k 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)nT⊆E∑{a}:T∑i=1∏k1−ynyai
也就是说，每多出一个连通块，额外产生的贡献就是 z = n y 1 − y z=\dfrac{ny}{1-y} z=1−yny。乘上每个连通块的大小再乘个常数就得到了答案。我们可以边 dp 边计算这个的贡献。
设 d p ( u , k ) dp(u,k) dp(u,k) 表示 u u u 为根的子树中，根所在连通块大小为 k k k，已经产生的总贡献。注意因为根还没有确定，所以根的贡献不计入。
可以写出方程：
d p ( u , k ) ← d p ( u , k ) ∑ i = 0 ∞ z i ⋅ d p ( v , i ) + ∑ i = 0 k d p ( u , i ) d p ( 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(u,k)←dp(u,k)i=0∑∞zi⋅dp(v,i)+i=0∑kdp(u,i)dp(v,k−i)
用树形 dp 的套路可以做到 O ( n 2 ) O(n^2) O(n2)，可是还不够。但注意到这东西又乘常数又卷积的，考虑生成函数。即设
d p u ( x ) = ∑ i = 0 ∞ d p ( u , i ) x i dp_u(x)=\sum_{i=0}^{\infin}dp(u,i)x^i dpu(x)=i=0∑∞dp(u,i)xi
对于这个 ∑ i ⋅ d p ( v , i ) \sum i\cdot dp(v,i) ∑i⋅dp(v,i)，发现它就是 d p u ′ ( 1 ) dp_u'(1) dpu′(1)（？）。我们设 f u = z ⋅ d p u ′ ( 1 ) f_u=z\cdot dp_u'(1) fu=z⋅dpu′(1)，发现答案就是 f 1 f_1 f1 乘上前面那坨常数。
方程可以改写成
d p u ( x ) ← d p u ( x ) f v + d p u ( x ) d p v ( x ) dp_u(x)\leftarrow dp_u(x)f_v+dp_u(x)dp_v(x) dpu(x)←dpu(x)fv+dpu(x)dpv(x)
为了推出 f u f_u fu，求个导。
d p u ′ ( x ) ← d p u ′ ( x ) f v + d p u ′ ( x ) d p v ( x ) + d p u ( x ) d p v ′ ( x ) dp_u'(x)\leftarrow dp_u'(x)f_v+dp_u'(x)dp_v(x)+dp_u(x)dp_v'(x) dpu′(x)←dpu′(x)fv+dpu′(x)dpv(x)+dpu(x)dpv′(x)
代入 x = 1 x=1 x=1,再乘个 z z z。
z ⋅ d p u ′ ( 1 ) ← z ⋅ d p u ′ ( 1 ) f v + z ⋅ d p u ′ ( 1 ) d p v ( 1 ) + d p u ( 1 ) ⋅ z ⋅ d p v ′ ( 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) z⋅dpu′(1)←z⋅dpu′(1)fv+z⋅dpu′(1)dpv(1)+dpu(1)⋅z⋅dpv′(1)
设 g u = d p u ( 1 ) g_u=dp_u(1) gu=dpu(1)，就得到了优美的方程
f u ← f u f v + f u g v + f v g u f_u\leftarrow f_uf_v+f_ug_v+f_vg_u fu←fufv+fugv+fvgu
g u g_u gu 用最开始的方程算
g u ← g u ( f v + g v ) g_u\leftarrow g_u(f_v+g_v) gu←gu(fv+gv)
就可以 O ( n ) O(n) O(n) 求了。
另一种等价的思路是设 f ( u , 0 / 1 ) f(u,0/1) f(u,0/1) 表示根结点所在的连通块是否计入了贡献，不过比求导还难想……
op=2
其实比前面还简单点……
用同样的思路推式子
∑ E 1 ∈ Tree ∑ E 2 ∈ Tree y n − ∣ E 1 ∩ E 2 ∣ \sum_{E_1\in \operatorname{Tree}}\sum_{E_2\in \operatorname{Tree}}y^{n-|E_1\cap E_2|} E1∈Tree∑E2∈Tree∑yn−∣E1∩E2∣
= ∑ S y n − ∣ S ∣ ∑ E 1 ∈ Tree ∑ E 2 ∈ Tree [ E 1 ∩ E 2 = S ] =\sum_{S}y^{n-|S|}\sum_{E_1\in \operatorname{Tree}}\sum_{E_2\in \operatorname{Tree}}[E_1\cap E_2=S] =S∑yn−∣S∣E1∈Tree∑E2∈Tree∑[E1∩E2=S]
= ∑ S y n − ∣ S ∣ ∑ S ⊆ T C 2 ( T ) ( − 1 ) ∣ T ∣ − ∣ S ∣ =\sum_{S}y^{n-|S|}\sum_{S\subseteq T}C^2(T)(-1)^{|T|-|S|} =S∑yn−∣S∣S⊆T∑C2(T)(−1)∣T∣−∣S∣
= ∑ T C 2 ( T ) ∑ S ⊆ T ( − 1 ) ∣ T ∣ − ∣ S ∣ y n − ∣ S ∣ =\sum_{T}C^2(T)\sum_{S\subseteq T}(-1)^{|T|-|S|}y^{n-|S|} =T∑C2(T)S⊆T∑(−1)∣T∣−∣S∣yn−∣S∣
= ∑ T C 2 ( T ) y n − ∣ T ∣ ∑ S ⊆ T ( − y ) ∣ T ∣ − ∣ S ∣ =\sum_{T}C^2(T)y^{n-|T|}\sum_{S\subseteq T}(-y)^{|T|-|S|} =T∑C2(T)yn−∣T∣S⊆T∑(−y)∣T∣−∣S∣
= ∑ T C 2 ( T ) y n − ∣ T ∣ ( 1 − y ) ∣ T ∣ =\sum_{T}C^2(T)y^{n-|T|}(1-y)^{|T|} =T∑C2(T)yn−∣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∏kai2yk(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)nT∑{a}:T∑i=1∏k1−yn2yai2
考虑一个大小为 s s s 的连通块，其贡献为后面那坨乘上自己内部的连边方案 s s − 2 s^{s-2} ss−2，即
n 2 y 1 − y s s \frac{n^2y}{1-y}s^s 1−yn2yss
构造出 EGF 然后 exp \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;
}