【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)\)；

    套用得：

    \[\begin{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_{S\subset 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\)种颜色，即：

    \[\begin{align}
    g(T) &= n^{m-2}\Pi_{i=1}^{m} a_i \\
    \Rightarrow \ ans &= y^n \sum_{T}z^{n-m} \times 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

  • 同理：
  \[\begin{align}
  ans &= y^n\sum_{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 \times i^{i-2} \times \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\n",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;
}