BZOJ5475 WC2019数树（prufer+容斥原理+树形dp+多项式exp）

　　因为一大堆式子实在懒得写题解了。首先用prufer推出CF917D用到的结论，然后具体见前言不搭后语的注释。

#include<iostream> 
#include<cstdio>
#include<cmath>
#include<cstdlib>
#include<cstring>
#include<algorithm>
#include<map>
using namespace std;
#define ll long long
#define P 998244353
#define N 600010
char getc(){char c=getchar();while ((c<'A'||c>'Z')&&(c<'a'||c>'z')&&(c<'0'||c>'9')) c=getchar();return c;}
int gcd(int n,int m){return m==0?n:gcd(m,n%m);}
int read()
{
	int x=0,f=1;char c=getchar();
	while (c<'0'||c>'9') {if (c=='-') f=-1;c=getchar();}
	while (c>='0'&&c<='9') x=(x<<1)+(x<<3)+(c^48),c=getchar();
	return x*f;
}
int n,m,op,w,f[N][2],p[N],t;
int a[N],b[N],c[N],d[N],A[N],B[N],fac[N],r[N];
struct data{int to,nxt;
}edge[N<<1];
void addedge(int x,int y){t++;edge[t].to=y,edge[t].nxt=p[x],p[x]=t;}
map<int,int> qaq[N];
int ksm(int a,int k)
{
	int s=1;
	for (;k;k>>=1,a=1ll*a*a%P) if (k&1) s=1ll*s*a%P;
	return s;
}
int inv(int a){return ksm(a,P-2);}
void inc(int &x,int y){x+=y;if (x>=P) x-=P;}
void dfs(int k,int from)
{
	f[k][0]=1;
	for (int i=p[k];i;i=edge[i].nxt)
	if (edge[i].to!=from)
	{
		dfs(edge[i].to,k);
		f[k][0]=1ll*f[k][0]*(f[edge[i].to][0]+f[edge[i].to][1])%P;
	}
	f[k][1]=1ll*f[k][0]*w%P;
	for (int i=p[k];i;i=edge[i].nxt)
	if (edge[i].to!=from) inc(f[k][1],1ll*f[k][0]*inv(f[edge[i].to][0]+f[edge[i].to][1])%P*f[edge[i].to][1]%P);
}
void DFT(int n,int *a,int g)
{
    for (int i=0;i<n;i++) r[i]=(r[i>>1]>>1)|(i&1)*(n>>1);
    for (int i=0;i<n;i++) if (i<r[i]) swap(a[i],a[r[i]]);
    for (int i=2;i<=n;i<<=1)
    {
        int wn=ksm(g,(P-1)/i);
        for (int j=0;j<n;j+=i)
        {
            int w=1;
            for (int k=j;k<j+(i>>1);k++,w=1ll*w*wn%P)
            {
                int x=a[k],y=1ll*w*a[k+(i>>1)]%P;
                a[k]=(x+y)%P,a[k+(i>>1)]=(x-y+P)%P;
            }
        }
    }
}
void IDFT(int *a,int n)
{
    DFT(n,a,inv(3));
    int u=inv(n);
    for (int i=0;i<n;i++) a[i]=1ll*a[i]*u%P;
}
void mul(int *a,int *b,int n)
{
    DFT(n,a,3),DFT(n,b,3);
    for (int i=0;i<n;i++) a[i]=1ll*a[i]*b[i]%P;
    IDFT(a,n);
}
void Inv(int *a,int *b,int n)
{
    if (n==1) {for (int i=0;i<t;i++) b[i]=0;b[0]=inv(a[0]);return;}
    Inv(a,b,n>>1);
    for (int i=0;i<n;i++) A[i]=a[i];
    for (int i=n;i<(n<<1);i++) A[i]=0;
    n<<=1;
    DFT(n,A,3),DFT(n,b,3);
    for (int i=0;i<n;i++) b[i]=1ll*b[i]*(P+2-1ll*A[i]*b[i]%P)%P;
    IDFT(b,n);
    n>>=1;
    for (int i=n;i<(n<<1);i++) b[i]=0;
}
void trans(int *a,int *b,int n){for (int i=0;i<n-1;i++) b[i]=1ll*a[i+1]*(i+1)%P;}
void dx(int *a,int *b,int n){b[0]=0;for (int i=1;i<n;i++) b[i]=1ll*a[i-1]*inv(i)%P;}
void Ln(int *a,int t)
{
    for (int i=0;i<t;i++) b[i]=c[i]=0;
    trans(a,c,t);
    Inv(a,b,t>>1);
    mul(c,b,t);
    dx(c,a,t);
}
void Exp(int *a,int *b,int n)
{
    if (n==1) {b[0]=1;return;}
    Exp(a,b,n>>1);
    for (int i=0;i<(n>>1);i++) B[i]=b[i];
    for (int i=(n>>1);i<n;i++) B[i]=0;
    Ln(B,n);
    for (int i=0;i<n;i++) B[i]=(a[i]-B[i]+P)%P;
    B[0]=(B[0]+1)%P;
    for (int i=n;i<(n<<1);i++) B[i]=0;
    mul(b,B,n<<1);
    for (int i=n;i<(n<<1);i++) b[i]=0;
}
int main()
{
#ifndef ONLINE_JUDGE
	freopen("bzoj5475.in","r",stdin);
	freopen("bzoj5475.out","w",stdout);
	const char LL[]="%I64d
";
#else
	const char LL[]="%lld
";
#endif
	n=read(),m=read(),op=read();
	if (op==0)
	{
		for (int i=1;i<n;i++)
		{
			int x=read(),y=read();
			qaq[x][y]=qaq[y][x]=1;
		}
		int cnt=0;
		for (int i=1;i<n;i++)
		{
			int x=read(),y=read();
			if (qaq[x][y]) cnt++;
		}
		cout<<ksm(m,n-cnt);
	}
	if (op==1)
	{
		if (m==1) {cout<<ksm(n,n-2);return 0;}
		for (int i=1;i<n;i++)
		{
			int x=read(),y=read();
			addedge(x,y),addedge(y,x);
		}
		w=1ll*n*inv(inv(m)-1)%P;
		dfs(1,1);
		cout<<1ll*ksm(m,n)*ksm(n,n-2)%P*ksm(inv(w),n)%P*f[1][1]%P;
		//f(i)钦定i条边相同时树的个数 
		//n^(n-i-2)*Σ∏size 
		//g(i)恰有i条边相同时树的个数
		//g(i)=Σ(-1)^(j-i)*C(j,i)*f(j) 
		//ans=Σg(i)*m^(n-i) 
		//ans=Σm^(n-i)*Σ(-1)^(j-i)*C(j,i)*f(j)
		//ans=Σf(i)*Σ(-1)^(i-j)*C(i,j)*m^(n-j)
		//ans=m^n*Σf(i)*Σ(-1)^(i-j)*C(i,j)*(1/m)^j
		//ans=m^n*Σf(i)*(1/m-1)^i
		//ans=m^n*n^(n-2)*Σ((1/m-1)/n)^i*Σ∏size 
		//p=n/(1/m-1)
		//ans=m^n*n^(n-2)*(1/p)^n*Σp^i*Σ∏size (i=1~n)
		//f[i][j] i子树 根所在连通块大小为j时 贡献之和
		//每个连通块选一个点 贡献之和 每选一个点乘p 
		//选一些边将其断开
		//f[i][0/1]i所在连通块是否选了点的贡献之和 
	}
	if (op==2)
	{
		if (m==1) {cout<<ksm(n,2*(n-2));return 0;}
		fac[0]=1;for (int i=1;i<=n;i++) fac[i]=1ll*fac[i-1]*i%P;
		w=inv(m)-1;
		for (int i=1;i<=n;i++) a[i]=1ll*n*n%P*inv(w)%P*ksm(i,i)%P*inv(fac[i])%P;
		int t=1;while (t<=(n+1<<1)) t<<=1;
		Exp(a,d,t);
		cout<<1ll*ksm(m,n)*ksm(w,n)%P*inv(ksm(n,4))%P*d[n]%P*fac[n]%P;
		//假设钦定了各连通块的点
		//m^n*Σ(n^(|S|-2)*(∏Sj^Sj)*(1/m-1)^(n-|S|)
		//m^n*(1/m-1)^n/n^4*(∏Sj^Sj*n^2*(1/m-1)^(-|S|)
		//m^n*(1/m-1)^n/n^4*(∏Sj^Sj*n^2/(1/m-1))
		//exp直接钦定 
	}
	return 0;
}