[WC2019]数树

XXXV.[WC2019]数树

首先将问题从”存在路径同时属于两个树“，先转换成被两个树上同时存在的边连成的连通块中的所有边须有相同颜色。进一步地，因为两棵树的并必然是森林，而森林的连通块数即为点数减边数，因此一对树 \(T_1,T_2\) 的贡献便是 \(f(T_1,T_2)=y^{n-|T_1\cap T_2|}\)。

则三问分别转换为：

0. \(f(T_1,T_2)\)
1. \(\sum\limits_{T_2}f(T_1,T_2)\)
2. \(\sum\limits_{T_1}\sum\limits_{T_2}f(T_1,T_2)\)

Situation 1. \(y=1\)

这种情形直接特判掉罢。

0. \(1\)
1. \(n^{n-2}\)
2. \(n^{2(n-2)}\)

Situation 2. \(op=0\)

\(T_1\cap T_2\) 可以直接求出，随便做就行了

Situation 3. \(op=1\)

考虑式子 \(\sum\limits_{T_2}f(T_1,T_2)\)。

我们换成枚举 \(S=T_1\cap T_2\)，得到 \(\sum\limits_{S}y^{n-|S|}\sum\limits_{T_2}[T_1\cap T_2=S]\)。

后面那个 \(=\) 不好处理，考虑容斥变为 \(S\subseteq T_1\cap T_2\)，也即 \(S\subseteq T_1\land S\subseteq T_2\)。则 \(y^{n-|S|}=\sum\limits_{S'\subseteq S}\sum\limits_{S''\subseteq S'}(-1)^{|S'|-|S''|}y^{n-|S''|}\)。

于是则整个式子变为 \(\sum\limits_{T_2}\sum\limits_{S\subseteq T_1\cap T_2}\sum\limits_{S'\subseteq S}(-1)^{|S|-|S'|y^{n-|S'|}}\)。交换枚举顺序可得 \(\sum\limits_{S\subseteq T_1}\sum\limits_{S'\subseteq S}(-1)^{|S|-|S'|}y^{n-|S'|}\sum\limits_{T_2\supseteq S}\)。

考虑将 \(|S|\) 与 \(|S'|\) 的项分开，就得到了 \(\sum\limits_{S\subseteq T_1}y^{n-|S|}\sum\limits_{S'\subseteq S}(-y)^{|S|-|S'|}\sum\limits_{T_2\supseteq S}\)。

考虑枚举 \(|S'|\) 为 \(k\)，得到 \(\sum\limits_{S\subseteq T_1}y^{n-|S|}\sum\limits_{k=0}^{|S|}\dbinom{|S|}{k}(-y)^{|S|-k}\sum\limits_{T_2\supseteq S}\)。直接上二项式定理，得到 \(\sum\limits_{S\subseteq T_1}y^{n-|S|}(1-y)^{|S|}\sum\limits_{T_2\supseteq S}\)。

下面考虑计算 \(\sum\limits_{T_2\supseteq S}\)。众所周知的是，一个点数为 \(n\)，连通块大小集合为 \(\{a\}\)，连通块数为 \(k\) 的森林，补成树的方案是 \(n^{k-2}\prod a_i\)。于是代入得到 \(\sum\limits_{S\subseteq T_1}y^{n-|S|}(1-y)^{|S|}n^{n-|S|-2}\prod\limits_{i=1}^{n-|S|}a_i\)。将所有指数上有 \(|S|\) 的东西摊入每项乘积里面，得到 \(\dfrac{(1-y)^n}{n^2}\sum\limits_{S\subseteq T_1}\prod\dfrac{ny}{(1-y)}a_i\)。这就相当于，树上一个大小为 \(a_i\) 的连通块贡献 \(Ka_i\)，其中 \(K\) 是常数。

直接DP的话，可以设 \(f_{i,j}\) 表示 \(i\) 子树中，所有位置 \(i\) 所在连通块大小为 \(j\) 的方案的答案之和。暴力转移即可树上背包做到 \(O(n^2)\)。

现在考虑 \(Ka_i\) 的组合意义：在连通块中任选一个点产生 \(K\) 的答案的方案数。于是设 \(f_{i,0/1}\) 表示 \(i\) 子树中，\(i\) 所在连通块未/已计算贡献，此时的答案之和，即可做到 \(O(n)\)。

Situation 4. \(op=2\)

考虑式子 \(\sum\limits_{T_1}\sum\limits_{T_2}f(T_1,T_2)\)。直接将我们上面的容斥套上去，得到 \(\dfrac{(1-y)^n}{n^2}\sum\limits_{T_1}\sum\limits_{S\subseteq T_1}\prod\dfrac{ny}{(1-y)}a_i\)。发现，此时我们可以将 \(\sum\limits_{T_1}\) 也扔到后面去，同样也是包含集合 \(S\) 的树数，即 \(n^{k-2}\prod a_i\)。于是就得到了 \(\dfrac{(1-y)^n}{n^4}\sum\limits_S\prod\dfrac{n^2y}{(1-y)}a_i^2\)。前面是常数，考虑后面。

因为这里不对 \(S\) 有除了是棵生成森林外的任何限制，所以我们可以考虑对 \(S\) 中每个连通块计算。考虑若固定了 \(S\) 的联通性（即固定每个连通块的大小，但不固定连边的情况）时的答案：每个大小为 \(a_i\) 的连通块，都可连成 \(a_i^{a_i-2}\) 种本质不同的生成树，所以一个大小为 \(a\) 的连通块的贡献即为 \(\dfrac{n^2y}{(1-y)}a^a\)。

于是问题转换为，单个连通块的生成函数是 \(f(x)=\sum\limits_{i=1}^n\dfrac{n^2y}{(1-y)}i^ix^{i}\)，求有标号的整张图的生成函数 \(g(x)\)。

这里有一个结论 \(\hat g=\exp\hat f\)，其中 \(\hat{}\) 是 EGF 的意思。其组合意义是，\(\exp\) 是带标号背包操作，\(\ln\) 是其反演。证明可以通过暴力泰勒展开 \(\exp\hat f\)，发现其与 \(g\) 的实际意义相同。

时间复杂度 \(O(n\log n)\)。

代码：

#include<bits/stdc++.h>
using namespace std;
const int mod=998244353;
int ksm(int x,int y=mod-2){int z=1;for(;y;y>>=1,x=1ll*x*x%mod)if(y&1)z=1ll*z*x%mod;return z;}
int n,m,op;
namespace Yequals1{
	void solve(){
		if(op==0)puts("1");
		if(op==1)printf("%d\n",ksm(n,n-2));
		if(op==2)printf("%d\n",ksm(n,(n-2)<<1));
		exit(0); 
	}
}
namespace Op0{
	map<pair<int,int>,bool>mp;
	void solve(){
		for(int i=1,x,y;i<n;i++){
			scanf("%d%d",&x,&y);
			if(x>y)swap(x,y);
			mp[make_pair(x,y)]=true;
		}
		int cnt=n;
		for(int i=1,x,y;i<n;i++){
			scanf("%d%d",&x,&y);
			if(x>y)swap(x,y);
			if(mp.find(make_pair(x,y))!=mp.end())cnt--;
		}
		printf("%d\n",ksm(m,cnt));
		exit(0);
	}
}
namespace Op1{
	vector<int>v[100100];
	int f[100100][2],K;
	void dfs(int x,int fa){
		f[x][0]=1,f[x][1]=K;
		for(auto y:v[x])if(y!=fa){
			dfs(y,x);
			f[x][1]=(1ll*f[x][1]*f[y][1]+1ll*f[x][1]*f[y][0]+1ll*f[x][0]*f[y][1])%mod;
			f[x][0]=(1ll*f[x][0]*f[y][0]+1ll*f[x][0]*f[y][1])%mod;
		}
	}
	void solve(){
		K=1ll*n*m%mod*ksm(1-m+mod)%mod;
		for(int i=1,x,y;i<n;i++)scanf("%d%d",&x,&y),v[x].push_back(y),v[y].push_back(x);
		dfs(1,0);
		printf("%d\n",1ll*f[1][1]*ksm(1-m+mod,n)%mod*ksm(n,mod-3)%mod);
		exit(0);
	}
}
namespace Op2{
	namespace Poly{
		const int N=1<<20;
		const int G=3;
		int A[N],B[N],C[N],D[N],rev[N],lim,invlim;
		void NTT(int*a,int tp){
			for(int i=0;i<lim;i++)if(rev[i]>i)swap(a[i],a[rev[i]]);
			for(int md=1;md<lim;md<<=1){
				int rt=ksm(G,(mod-1)/(md<<1));if(tp==-1)rt=ksm(rt);
				for(int stp=md<<1,pos=0;pos<lim;pos+=stp)for(int i=0,w=1;i<md;i++,w=1ll*w*rt%mod){
					int x=a[pos+i],y=1ll*a[pos+md+i]*w%mod;
					a[pos+i]=(x+y)%mod,a[pos+md+i]=(x+mod-y)%mod;
				} 
			}
			if(tp==-1)for(int i=0;i<lim;i++)a[i]=1ll*a[i]*invlim%mod;
		} 
		void mul(int*a,int*b,int*c,int LG){
			invlim=ksm(lim=1<<LG);for(int i=0;i<lim;i++)rev[i]=(rev[i>>1]>>1)|((i&1)<<(LG-1));
			for(int i=0;i<lim;i++)A[i]=B[i]=0;
			for(int i=0;i<(lim>>1);i++)A[i]=a[i],B[i]=b[i];
//			for(int i=0;i<(lim>>1);i++)printf("%d ",a[i]);putchar('*');
//			for(int i=0;i<(lim>>1);i++)printf("%d ",b[i]);putchar('=');
			NTT(A,1),NTT(B,1);
			for(int i=0;i<lim;i++)A[i]=1ll*A[i]*B[i]%mod;
			NTT(A,-1);
			for(int i=0;i<lim;i++)c[i]=A[i];
//			for(int i=0;i<lim;i++)printf("%d ",c[i]);putchar('\n');
		}
		void inv(int *a,int *b,int LG){
			b[0]=ksm(a[0]);
			for(int k=1;k<=LG+1;k++){
				mul(b,a,C,k);
				for(int i=0;i<(1<<k);i++)C[i]=(mod-C[i])%mod;
				(C[0]+=2)%=mod;
				mul(C,b,b,k);
			}
		}
		void diff(int *a,int *b,int len){for(int i=0;i<len;i++)b[i]=1ll*a[i+1]*(i+1)%mod;b[len-1]=0;}
		void inte(int *a,int *b,int len){for(int i=len-1;i;i--)b[i]=1ll*a[i-1]*ksm(i)%mod;b[0]=0;}
		void ln(int *a,int *b,int LG){
			inv(a,b,LG);
			diff(a,C,1<<LG);
			mul(b,C,b,LG+1);
			inte(b,b,1<<LG);
		}
		void exp(int *a,int *b,int LG){
			b[0]=1;
			for(int k=1;k<=LG+1;k++){
				ln(b,D,k-1);
				for(int i=0;i<(1<<(k-1));i++)D[i]=(a[i]-D[i]+mod)%mod;
				D[0]=(D[0]+1)%mod;
				mul(b,D,b,k);
			}
		}
	}
	using namespace Poly;
	int fac[N],caf[N],f[N],g[N],LG;
	void solve(){
		fac[0]=1;for(int i=1;i<=n;i++)fac[i]=1ll*fac[i-1]*i%mod;
		caf[n]=ksm(fac[n]);for(int i=n;i;i--)caf[i-1]=1ll*caf[i]*i%mod;
		int K=1ll*n*n%mod*m%mod*ksm(1-m+mod)%mod;
		for(int i=1;i<=n;i++)f[i]=1ll*ksm(i,i)*K%mod*caf[i]%mod;
//		for(int i=1;i<=n;i++)printf("%d ",f[i]);puts("");
		while((1<<LG)<=n)LG++;
		exp(f,g,LG);
//		for(int i=1;i<=n;i++)printf("%d ",g[i]);puts("");
		printf("%d\n",1ll*g[n]*fac[n]%mod*ksm(1+mod-m,n)%mod*ksm(n,mod-5)%mod);
		exit(0);
	}
}
int main(){
	scanf("%d%d%d",&n,&m,&op);
	if(m==1)Yequals1::solve();
	if(op==0)Op0::solve();
	if(op==1)Op1::solve();
	if(op==2)Op2::solve();
	return 0;
}