[WC2019]数树

Description

有两棵n个点的树，T1和T2，每个点可以填一个[1,y]的颜色
若两棵树有一条边重合这条边两个端点的颜色必须相同
有三种问题：
op=0，给出T1和T2问答案
op=1，给出T1问所有T2的答案和
op=2，给出n问所有的T1和T2的答案和
n<=100000

Solution

op=0大家都会
op=1的话，考虑算恰好有i条重边的$y^{-i}$的和
考虑到$z^m={\sum }_{i=0}^m(z-1{)}^i{(}_i^m)$
于是我们可以变成，枚举i条重边，其余边随意的(z-1)^i的和
i条边会把原树分成n-i个连通块，设ai为连通块大小，根据purfer序我们可以推出生成树个数为$n^{n-i-2}\prod ai$
考虑到$\prod ai$的组合意义为从每个连通块里选一个点的方案数
设F[x][0…1]表示当前x为根的连通块有没有选这个点
直接Dp是O(n)的
op=2根据op=1断掉每条重边求出一个连通块的EGF然后exp一下就好了

Code

#include <map>
#include <cstdio>
#include <cstring>
#include <algorithm>
#define fo(i,a,b) for(int i=a;i<=b;i++)
#define fd(i,a,b) for(int i=a;i>=b;i--)
#define rep(i,a) for(int i=lst[a];i;i=nxt[i])
using namespace std;

typedef long long ll;

int read() {
	char ch;
	for(ch=getchar();ch<'0'||ch>'9';ch=getchar());
	int x=ch-'0';
	for(ch=getchar();ch>='0'&&ch<='9';ch=getchar()) x=x*10+ch-'0';
	return x;
}

const int N=1e6+5,Mo=998244353;

int pwr(int x,int y) {
	int z=1;
	for(;y;y>>=1,x=(ll)x*x%Mo)
		if (y&1) z=(ll)z*x%Mo;
	return z;
}

map<int,int> e[N];

int n,y,z;

void solve_0() {
	fo(i,1,n-1) {
		int x=read(),y=read();
		e[x][y]=e[y][x]=1;
	}
	int cnt=0;
	fo(i,1,n-1) {
		int x=read(),y=read();
		cnt+=e[x][y];
	}
	printf("%d\n",pwr(y,n-cnt));
}

int t[N],nxt[N],lst[N],l;
void add(int x,int y) {t[++l]=y;nxt[l]=lst[x];lst[x]=l;}

int dp[N][2];

void Dp(int x,int y) {
	rep(i,x) if (t[i]!=y) Dp(t[i],x);
	dp[x][0]=dp[x][1]=1;int s0,s1;
	rep(i,x)
		if (t[i]!=y) {
			// 选
			s0=(ll)dp[x][0]*dp[t[i]][1]%Mo*n%Mo;
			s1=(ll)dp[x][1]*dp[t[i]][1]%Mo*n%Mo;
			// 不选
			(s0+=(ll)dp[x][0]*dp[t[i]][0]%Mo*z%Mo)%=Mo;
			(s1+=(ll)dp[x][1]*dp[t[i]][0]%Mo*z%Mo)%=Mo;
			(s1+=(ll)dp[x][0]*dp[t[i]][1]%Mo*z%Mo)%=Mo;
			dp[x][0]=s0;dp[x][1]=s1;
		}
}

void solve_1() {
	fo(i,1,n-1) {
		int x=read(),y=read();
		add(x,y);add(y,x);
	}
	if (y==1) {
		printf("%d\n",pwr(n,n-2));
		return;
	}
	z=pwr(y,Mo-2);z--;
	Dp(1,0);
	int ans=(ll)dp[1][1]*pwr(n,Mo-2)%Mo;
	ans=(ll)ans*pwr(y,n)%Mo;
	printf("%d\n",ans);
}

ll W[2][N];

void init(int b) {
	for(int i=1;i<(1<<b);i<<=1){
		ll wn=pwr(3,(Mo-1)/(i<<1));
		for(int j=0;j<i;++j) W[1][i+j]=(j?wn*W[1][i+j-1]%Mo:1);
		wn=pwr(3,Mo-1-(Mo-1)/(i<<1));
		for(int j=0;j<i;++j) W[0][i+j]=(j?wn*W[0][i+j-1]%Mo:1);
	}
}

void DFT(ll *a,int len,int flag) {
	if (flag==-1) flag=0;
	for(int i=0,j=0;i<len;++i){
		if (i<j) swap(a[i],a[j]);
		for(int k=len>>1;(j^=k)<k;k>>=1);
	}
	for(int i=1;i<len;i<<=1)
		for(int j=0;j<len;j+=(i<<1))
			for(int k=0;k<i;++k) {
				ll x=a[j+k],y=a[j+k+i]*W[flag][i+k]%Mo;
				a[j+k]=(x+y)%Mo;
				a[j+k+i]=(x-y)%Mo;
			}
	ll inv=pwr(len,Mo-2);
	if (!flag) for(int i=0;i<len;i++) a[i]=a[i]*inv%Mo;
}

ll c[N];

void get_Inv(ll *a,ll *b,int n) {
	if (n==1) {b[0]=pwr(a[0],Mo-2);return;}
	get_Inv(a,b,n>>1);
	int len=n<<1;
	fo(i,0,n-1) c[i]=a[i];fo(i,n,len-1) c[i]=0;
	fo(i,(n>>1),len-1) b[i]=0;
	DFT(c,len,1);DFT(b,len,1);
	fo(i,0,len-1) b[i]=(2*b[i]-b[i]*b[i]%Mo*c[i])%Mo;
	DFT(b,len,-1);
	fo(i,n,len-1) b[i]=0;
}

ll f[N],g[N];

void get_ln(ll *a,ll *b,int n) {
	fo(i,0,n-2) f[i]=a[i+1]*(i+1)%Mo;f[n-1]=0;
	get_Inv(a,g,n);
	int len=n<<1;
	fo(i,n,len-1) f[i]=g[i]=0;
	DFT(f,len,1);DFT(g,len,1);
	fo(i,0,len-1) f[i]=f[i]*g[i]%Mo;
	DFT(f,len,-1);
	fo(i,1,n-1) b[i]=f[i-1]*pwr(i,Mo-2)%Mo;
	b[0]=0;
}

ll h[N];

void get_exp(ll *a,ll *b,int n) {
	if (n==1) {b[0]=1;return;}
	get_exp(a,b,n>>1);
	get_ln(b,h,n);
	fo(i,0,n-1) h[i]=(a[i]-h[i]+Mo)%Mo;
	(h[0]=h[0]+1)%=Mo;
	int len=n<<1;
	fo(i,n,len-1) h[i]=b[i]=0;
	DFT(h,len,1);DFT(b,len,1);
	fo(i,0,len-1) b[i]=b[i]*h[i]%Mo;
	DFT(b,len,-1);
	fo(i,n,len-1) b[i]=0;
}

ll F[N],G[N],fac[N],inv[N];

void solve_2() {
	init(18);
	if (y==1) {
		printf("%d\n",pwr(pwr(n,n-2),2));
		return;
	}
	z=pwr(y,Mo-2);int iz=pwr(z-1,Mo-2);
	fac[0]=1;fo(i,1,n) fac[i]=fac[i-1]*i%Mo;
	inv[n]=pwr(fac[n],Mo-2);fd(i,n-1,0) inv[i]=(ll)inv[i+1]*(i+1)%Mo;
	fo(i,1,n) F[i]=(ll)n*n%Mo*iz%Mo*pwr(i,i)%Mo*inv[i]%Mo;
	int len=1;for(;len<=n;len<<=1);get_exp(F,G,len);
	int ans=(ll)G[n]*fac[n]%Mo;
	ans=(ll)ans*pwr(pwr(n,Mo-2),4)%Mo*pwr(z-1,n)%Mo;
	ans=(ll)ans*pwr(y,n)%Mo;
	printf("%d\n",(ans+Mo)%Mo);
}

int main() {
	freopen("tree.in","r",stdin);
	freopen("tree.out","w",stdout);
	n=read();y=read();int op=read();
	if (op==0) solve_0();
	if (op==1) solve_1();
	if (op==2) solve_2();
	return 0;
}