【loj2983】【WC2019】数树

题目

两颗(n)个点的树T1和T2，有(y)种颜色;

现在给每个点染色，要求公共边端点的颜色相同，求：

​ 1.op=0 , T1和T2都确定，求合法染色方案数；

​ 2.op=1 , T1确定，求所有T2的合法染色方案数的和 ;

​ 3.op=2 , 求所有(T1,T2)合法染色方案数的和；

(mod 998244353)的值；

$1 le n le 10^5 , 1 le y lt 998244353 , op = { 0,1,2 } $ ;

题解

• op=0

• map即可

• op=1

• Part 1

  • 子集反演:

    [g(S) = sum_{S subset T} f(T) Leftrightarrow f(S) = sum_{S subset T} (-1)^{|T|-|S|}f(T) \ ]

    设(T_1 cap T_2 = S)的方案数为(f(S))，(T_1 cap T_2 supset S)的方案数为(g(S))；

    套用得：

    [egin{align} ans &= sum_{S} f(S) y^{n-|S|}\ &= y^n sum_{S} y^{-|S|} sum_{S subset T}(-1)^{|T|-|S|} g(T) \ &= y^n sum_{T} g(T) sum_{Ssubset T}y^{-|S|}(-1)^{|T|-|S|} \ &= y^n sum_{T} g(T) sum_{i=0}^{|T|} (^{|T|}_{ i})(frac{1}{y})^{i}(-1)^{|T|-i}\ &= y^n sum_{T} g(T) (frac{1}{y}-1)^{|T|} \ 令z&= frac{1}{y}-1\ ans &= y^n sum_{T} z^{|T|}g(T) \ end{align} ]

• Part 2

  • prufer序列求(g(T))：

    (T)会形成(m = n-|T|)个连通块，设大小为(a_i)，(g(T))相当于将其重新连成树的方案数：

    [g(T) = sum_{sum d_i = m-2} (m-2)! Pi_{i=1}^{m}frac{a_i^{d_i+1}}{d_i!} ]

    考虑先构造一个prufer序列，再在最后依次加入(1-m)得到一个长度为(2m-2)的度数序列；

    (g(T))的组合意义就是对这些点染色，(i)号连通块有(a_i)种颜色可以染；

    对前m-2个位置有n中不同的颜色，后m个位置每个有(a_i)种颜色，即：

    [egin{align} g(T) &= n^{m-2}Pi_{i=1}^{m} a_i \ Rightarrow ans &= y^n sum_{T}z^{n-m} imes n^{m-2}Pi_{i=1}^{m} a_i\ ans &= frac{y^n z^n}{n^2} sum_{T}z^{-m}n^{m}Pi_{i=1}^{m} a_i end{align} ]

    把(Pi_{i=1}^{m} a_i)看成每个连通块选一个点，(f_{ { u,0/1 } })表示是否选了的贡献即可(dp)；

• op=2

• Part 1

  • 同理：
  [egin{align} ans &= y^nsum_{T} z^{|T|} g^2(T) \ &= frac{y^n z^n}{n^4} sum_{T}z^{-m}n^{2m}Pi_{i=1}^{m} a_i^2 \ end{align} ]

• (sum_{T})后面那一坨可以看成是(m)个联通块组成的图；

• 连通块的EGF为(F(x) = frac{n^2}{z} i^2 imes i^{i-2} imes frac{x^i}{i!} Rightarrow frac{n^2i^i}{zi!}x^i)

• 所以图的EGF为$G(x) = e^{F(x)} Rightarrow ans = frac{y^n z^n n!}{n^4} [x^n]G(x) $

• (这个策爷的论文里有讲)

#include<bits/stdc++.h>
#define ll long long 
#define mk make_pair
#define mod 998244353
using namespace std;
const int N=100010;
int n,Y,Z,op;
char gc(){
	static char*p1,*p2,s[1000000];
	if(p1==p2)p2=(p1=s)+fread(s,1,1000000,stdin);
	return(p1==p2)?EOF:*p1++;
}
int rd(){
	int x=0;char c=gc();
	while(c<'0'||c>'9')c=gc();
	while(c>='0'&&c<='9')x=(x<<1)+(x<<3)+c-'0',c=gc();
	return x;
}
int pw(int x,int y){
	int re=1;
	if(y<0)y+=mod-1;
	while(y){
		if(y&1)re=(ll)re*x%mod;
		y>>=1;x=(ll)x*x%mod;
	}
	return re;
}
namespace subtask0{
	typedef pair<int,int>pii;
	map<pii,bool>mp;
	void solve(){
		for(int i=1;i<n;++i){
			int u=rd(),v=rd();
			if(u>v)swap(u,v);
			mp[mk(u,v)]=1;
		}
		int cnt=0;
		for(int i=1;i<n;++i){
			int u=rd(),v=rd();
			if(u>v)swap(u,v);
			if(mp[mk(u,v)])cnt++;
		}
		printf("%d
",pw(Y,n-cnt));
	}
}

namespace subtask1{
	int f[N][2],o=1,hd[N],pre,ipre;
	struct Edge{int v,nt;}E[N<<1];
	void adde(int u,int v){
		E[o]=(Edge){v,hd[u]};hd[u]=o++;
		E[o]=(Edge){u,hd[v]};hd[v]=o++;
	}
	void dfs(int u,int fa){
		f[u][0]=f[u][1]=pre;
		for(int i=hd[u];i;i=E[i].nt){
			int v=E[i].v;
			if(v==fa)continue;
			dfs(v,u);
			int t0 = f[u][0] , t1 = f[u][1]; 
			f[u][0] = (1ll * t0 * f[v][1] %mod + 1ll * t0 * f[v][0] %mod * ipre %mod) %mod;
			f[u][1] = (1ll * t1 * f[v][1] %mod + 1ll * (1ll*t0*f[v][1]+1ll*t1*f[v][0]) %mod * ipre %mod) %mod;
		}
	}
	void solve(){
		if(!Z){cout<<pw(n,n-2)<<endl;return;}
		for(int i=1;i<n;++i)adde(rd(),rd());
		pre = 1ll * n * pw(Z , mod-2) %mod; 
		//pre = 1;
		ipre = pw(pre , mod-2); 
		dfs(1,0);
		int ans = 1ll * f[1][1] * pw(1ll*Y*Z%mod , n) %mod * pw(1ll*n*n%mod , mod-2) %mod;
		cout << ans << endl;
	}
}
namespace subtask2{
	const int N2=N<<2;
	int fac[N2],ifac[N2],iv[N2],a[N2],b[N2],L,G=3,rev[N2];
	void ntt(int*A,int len,int f){
		for(L=0;(1<<L)<len;++L);
		for(int i=0;i<len;++i){
			rev[i]=(rev[i>>1]>>1)|((i&1)<<(L-1));
			if(i<rev[i])swap(A[i],A[rev[i]]);
		}
		for(int i=1;i<len;i<<=1){
			int wn=pw(G , f * (mod-1)/i/2);
			for(int j=0;j<len;j+=i<<1){
				int w=1;
				for(int k=0;k<i;++k,w=1ll*w*wn%mod){
					int x=A[j+k],y=1ll*w*A[j+k+i]%mod; 
					A[j+k]=(x+y)%mod,A[j+k+i]=(x-y+mod)%mod;
				}
			}
		}
		if(!~f)for(int i=0;i<len;++i)A[i]=(ll)A[i]*iv[len]%mod;
	}
	void cls(int*A,int l,int r){for(int i=l;i<r;++i)A[i]=0;}
	void cpy(int*A,int*B,int l){for(int i=0;i<l;++i)A[i]=B[i];}
	void der(int*A,int l){for(int i=0;i<l-1;++i)A[i]=1ll*A[i+1]*(i+1)%mod;A[l-1]=0;}
	void dif(int*A,int l){for(int i=l-1;i;--i)A[i]=1ll*A[i-1]*iv[i]%mod;A[0]=0;}
	void inv(int*A,int*B,int l){
		if(l==1){B[0]=1;return;}
		static int t[N2];
		int len=l<<1;
		inv(A,B,l>>1);cls(B,l,len);
		cpy(t,A,l);cls(t,l,len);
		ntt(B,len,1);ntt(t,len,1);
		for(int i=0;i<len;++i)B[i] = B[i] * (2 - 1ll * t[i] * B[i] %mod +mod) %mod;
		ntt(B,len,-1);cls(B,l,len);
	}
	void ln(int*A,int*B,int l){
		static int t[N2];
		int len=l<<1;
		inv(A,B,l);
		cpy(t,A,l);cls(t,l,len);der(t,l);
		ntt(B,len,1);ntt(t,len,1);
		for(int i=0;i<len;++i)B[i] = 1ll * B[i] * t[i] %mod;
		ntt(B,len,-1);cls(B,l,len);
		dif(B,l);
	}
	void exp(int*A,int*B,int l){
		if(l==1){B[0]=1;return;}
		static int t[N2];
		int len=l<<1;
		exp(A,B,l>>1);cls(B,l,len);
		ln(B,t,l);
		for(int i=0;i<l;++i)t[i] = ( -t[i] + A[i] + mod)%mod;
		t[0]++;//!!!
		ntt(B,len,1);ntt(t,len,1);
		for(int i=0;i<len;++i)B[i] = 1ll * B[i] * t[i] %mod;
		ntt(B,len,-1);cls(B,l,len);
	}
	void solve(){
		if(!Z){cout<<pw(n,2*n-4)<<endl;return ;}
		int len=1;for(;len<=2*n;len<<=1);
		iv[1]=1;for(int i=2;i<=len;++i)iv[i]=1ll*(mod-mod/i)*iv[mod%i]%mod;
		for(int i=fac[0]=ifac[0]=1;i<=len;++i){
			ifac[i]=(ll)ifac[i-1]*iv[i]%mod;
			fac[i]=(ll)fac[i-1]*i%mod;		
		}
		int tmp = 1ll * n * n %mod * pw(Z,mod-2) %mod;
		for(int i=1;i<=n;++i)a[i] = 1ll * tmp * pw(i,i) %mod * ifac[i] %mod ;
		exp(a, b, len>>1 ); 
		int ans = 1ll * pw(1ll*Y*Z%mod , n) * pw((ll)n*n%mod*n%mod*n%mod , mod-2) %mod * fac[n] %mod * b[n] %mod;
		cout << ans << endl;
	}
}
int main(){
	freopen("tree.in","r",stdin);
	freopen("tree.out","w",stdout);
	n=rd();Y=rd();Z=pw(Y,mod-2)-1;op=rd();
	if(op==0)subtask0::solve();
	else if(op==1)subtask1::solve();
	else subtask2::solve();
	return 0;
}