[LibreOJ #2983]【WC2019】数树【计数】【DP】【多项式】

Description

此题含有三个子问题
问题1：
给出n个点的两棵树，记m为只保留同时在两棵树中的边时连通块的个数，求\(y^m\)
问题2：
给出n个点的一棵树，另外一棵树任意生成，求所有方案总的\(y^m\)的和
问题3：
两棵树均任意生成，求所有方案总的\(y^m\)的和
n<=100000,答案对998244353取模

Solution

问题1：

求出同时在两棵树中的边数v，容易得到m=n-v，直接算即可。

问题2：

记\(z=y^{-1}\)即y在模意义下的逆元，我们就是要求\(z^v\)的总和
直接计算比较困难
考虑这样一个转化\(z^v=\sum\limits_{i=0}^{v}(z-1)^i{v\choose i}\)
考虑它的组合意义，实际上就是我们任意枚举一个重合边集E的子集E0，贡献就是\((z-1)^{|E_0|}\)

我们不妨枚举原树边集的子集\(E_0\)，再计算有多少个生成树包含\(E_0\)
假设\(E_0\)构成了m个连通块，它们的点数分别是\(a_1,a_2,...,a_m\)

那么包含\(E_0\)的生成树个数就是\(n^{m-2}\prod\limits_{i=1}^{m} a_i\)，总的贡献就是\((z-1)^{n-m}n^{m-2}\prod\limits_{i=1}^{m} a_i\)
这东西考虑用prufer序来理解，我们将一个连通块看做一个大的点，那么m-2个连通块在prufer序上都有一个“父亲”，但是这个父亲是一个单点，因此是\(n^{m-2}\)，此外后面的大小之积相当于枚举每个连通块prufer序连出去的边是哪一条。

这显然可以用一个n^2的树形DP来完成，考虑优化
考虑上面的式子的组合意义，我们把常数因子提出来（\((z-1)^nn^{-2}\)），相当于每个连通块都有一个固定的贡献\((z-1)^{-1}n\)，此外再在每个连通块中选择一个点，求总的贡献和。

这就可以用树形DP来做了，记\(F[i][0/1]\)表示当前做完了i的子树，i所在的连通块是否已经贡献过。
直接讨论相邻边是否出现即可转移，时间复杂度\(O(N)\)，非常的巧妙。

问题3：

有了问题2的基础，我们来考虑问题3
不妨同样枚举重合边集的一个子集\(E_0\)，m,a的定义同上
那么贡献变成了\[(z-1)^{n-m}\left(n^{m-2}\prod a_i \right)^2\]
平方乘到指数上，依然是将常数因子提出来，把贡献分配到每个连通块上
此外由于两棵树都是生成的，因此连通块内部的边也不确定，还要乘上连通块内部的生成树个数

那么一个连通块i的贡献就是
\((z-1)^{-1}n^2a_i^2a_i^{a_i-2}=(z-1)^{-1}n^2a_i^{a_i}\)

那么上面的式子相当于连通块的带标号重复拼接
一个连通块的EGF就是\[F(x)=\sum\limits_{i>0}{(z-1)^{-1}n^2i^i\over i!}\]
拼接以后就是\([x^n]e^{F(x)}\)，乘上前面提出来的贡献就是答案了。

Code

#include <bits/stdc++.h>
#define fo(i,a,b) for(int i=a;i<=b;++i)
#define fod(i,a,b) for(int i=a;i>=b;--i)
#define N 100005
#define LL long long
#define mo 998244353
using namespace std;
int n,tp;
LL m;
LL ksm(LL k,LL n)
{
    LL s=1;
    for(;n;n>>=1,k=k*k%mo) if(n&1) s=s*k%mo;
    return s;
}
namespace subtask1
{
    map<int,bool> mp[N];
    void solve1()
    {
        fo(i,1,n-1)
        {
            int x,y;
            scanf("%d%d",&x,&y);
            mp[x][y]=mp[y][x]=1;
        }
        int c=0;
        fo(i,1,n-1)
        {
            int x,y;
            scanf("%d%d",&x,&y);
            if(mp[x][y]) c++;
        }
        printf("%lld\n",ksm(m,n-c));
    }
}
using namespace subtask1;

namespace subtask2
{
    LL f[N][2],rm,pm;
    int fs[N],m1,nt[2*N],dt[2*N];
    void link(int x,int y)
    {
        nt[++m1]=fs[x];
        dt[fs[x]=m1]=y;
    }
    void dp(int k,int fa)
    {
        f[k][0]=1,f[k][1]=0;
        for(int i=fs[k];i;i=nt[i])
        {
            int p=dt[i];
            if(p!=fa)
            {
                dp(p,k);
                f[k][1]=(f[k][1]*(f[p][0]+f[p][1])%mo+f[k][0]*f[p][1])%mo;
                f[k][0]=f[k][0]*(f[p][0]+f[p][1])%mo;
            }
        }
        f[k][1]=(f[k][1]+f[k][0]*pm%mo*(LL)n)%mo;
    }
    void solve2()
    {
        if(m==1) 
        {
            printf("%lld\n",ksm(n,n-2));
            return;
        }
        fo(i,1,n-1)
        {
            int x,y;
            scanf("%d%d",&x,&y);
            link(x,y),link(y,x);
        }   
        rm=ksm(m,mo-2);
        pm=ksm(rm-1,mo-2);
        dp(1,0);
        printf("%lld\n",f[1][1]*ksm(ksm(n,mo-2),2)%mo*ksm(rm-1,n)%mo*ksm(m,n)%mo);
    }
}
using namespace subtask2;

namespace subtask3
{
    #define M 262144
    LL wg[M+1],wi[M+1],a[M+1],b[M+1],js[M+1],ny[M+1],ns[M+1];
    int bit[M+1],cf[19],l2[M+1];
    void prp(int num)
    {
        fo(i,0,num)
        {
            wi[i]=wg[i*(M/num)];
            bit[i]=(bit[i>>1]>>1)|((i&1)<<(l2[num]-1));
        }
    }
    void NTT(LL *a,bool pd,int num)
    {
        LL v;
        fo(i,0,num-1) if(i<bit[i]) swap(a[i],a[bit[i]]);
        for(int h=1,m=2,l=num>>1;m<=num;h=m,m<<=1,l>>=1)
        {
            int c=(!pd)?l:-l;
            for(int j=0;j<num;j+=m)
            {
                LL *x=a+j,*y=a+h+j,*w=(!pd)?wi:wi+num;
                fo(i,0,h-1)
                {
                    v=*y * *w%mo;
                    *y=(*x-v+mo)%mo,*x=(*x+v)%mo;
                    x++,y++,w+=c;
                }   
            }
        }
        if(pd) fo(i,0,num-1) a[i]=a[i]*ny[num]%mo;
    }
    void getinv(int n,LL *a,LL *b)
    {
        static LL u1[M+1],u2[M+1];
        fo(i,0,cf[l2[n]]-1) b[i]=0;
        b[0]=ksm(a[0],mo-2);
        for(int m=1,t=2,num=4;m<n;m=t,t=num,num<<=1)
        {
            prp(num);
            fo(i,0,num-1) u1[i]=u2[i]=0;
            fo(i,0,m-1) u1[i]=b[i];
            fo(i,0,t-1) u2[i]=a[i];
            NTT(u1,0,num),NTT(u2,0,num);
            fo(i,0,num-1) u1[i]=u1[i]*u1[i]%mo*u2[i]%mo;
            NTT(u1,1,num);
            fo(i,0,t-1) b[i]=((LL)2*b[i]-u1[i]+mo)%mo;
        }
    }
    void getln(int n,LL *a,LL *b)
    {
        static LL u1[M+1],u2[M+1];
        int num=cf[l2[2*n+1]];
        getinv(n+1,a,u1);
        fo(i,n+1,num-1) u1[i]=0;
        prp(num);
        fo(i,0,n-1) u2[i]=a[i+1]*(LL)(i+1)%mo;
        fo(i,n,num-1) u2[i]=0;
        NTT(u1,0,num),NTT(u2,0,num);
        fo(i,0,num-1) u1[i]=u1[i]*u2[i]%mo;
        NTT(u1,1,num);
        b[0]=0;
        fo(i,1,n) b[i]=u1[i-1]*ny[i]%mo;
    }
    void getexp(int n,LL *a,LL *b)
    {
        static LL u1[M+1],u2[M+1];
        b[0]=1;
        fo(i,1,cf[l2[n]]) b[i]=0;
        for(int m=1,t=2,num=4;m<=n;m=t,t=num,num<<=1)
        {
            getln(t-1,b,u1);
            fo(i,0,t-1) u1[i]=(-u1[i]+a[i]+mo+mo)%mo;
            fo(i,t,num-1) u1[i]=0;
            u1[0]++;
            fo(i,0,m-1) u2[i]=b[i];
            fo(i,m,num-1) u2[i]=0;
            prp(num); 
            NTT(u1,0,num),NTT(u2,0,num);
            fo(i,0,num-1) u1[i]=u1[i]*u2[i]%mo;
            NTT(u1,1,num);
            fo(i,0,t-1) b[i]=u1[i];
        }
    }
    void solve3()
    {
        if(m==1) {printf("%lld\n",ksm(n,2*n-4));return;}
        wg[0]=1,wg[1]=ksm(3,(mo-1)/M);
        fo(i,2,M) wg[i]=wg[i-1]*wg[1]%mo;
        fo(i,0,18) l2[cf[i]=1<<i]=i;
        fod(i,M-1,2) if(!l2[i]) l2[i]=l2[i+1];
        
        js[0]=ns[0]=js[1]=ns[1]=ny[1]=1;
        fo(i,2,M) js[i]=js[i-1]*(LL)i%mo,ny[i]=(-ny[mo%i]*(LL)(mo/i)%mo+mo)%mo;
        fo(i,2,M) ns[i]=ns[i-1]*ny[i]%mo;
        
        LL rm=ksm(m,mo-2),pm=ksm(rm-1,mo-2);
        fo(i,1,n) a[i]=pm*(LL)n%mo*(LL)n%mo*ksm(i,i)%mo*ns[i]%mo;
        getexp(n,a,b);
        printf("%lld\n",b[n]*ksm(ksm(n,mo-2),4)%mo*ksm(rm-1,n)%mo*ksm(m,n)%mo*js[n]%mo);
    }
}
using namespace subtask3;

int main()
{
    cin>>n>>m>>tp;
    if(tp==0) solve1();
    else if(tp==1) solve2();
    else solve3();
}