Codeforces997D Cycles in product 【FFT】【树形DP】

题目大意：

给两个树，求环的个数。

题目分析：

出题人摆错题号系列。

通过画图很容易就能想到把新图拆在两个树上，在树上游走成环。

考虑DP状态F，G，T。F表示最终答案，T表示儿子不考虑父亲，G表示父亲不考虑儿子。T通过从下往上做NTT，G通过从上往下做NTT。F顺便做NTT。

最后做一下拼接就行。

代码：

#include<bits/stdc++.h>
using namespace std;

const int maxn = 4020;
const int mod = 998244353;
const int gg = 3;

int n[2],k;

vector <int> g[2][maxn];

int f[2][maxn][80],gi[2][maxn][80],T[2][maxn][80];

int C[100][100];

int fast_pow(int now,int pw){
    int ans = 1,bit = 1,dt = now;
    while(bit <= pw){
        if(bit & pw){ans = (1ll*ans*dt)%mod;}
        bit <<=1; dt = (1ll*dt*dt)%mod;
    }
    return ans;
}

void read(){
    scanf("%d%d%d",&n[0],&n[1],&k);
    for(int i=1;i<n[0];i++){
        int u,v; scanf("%d%d",&u,&v);
        g[0][u].push_back(v); g[0][v].push_back(u);
    }
    for(int i=1;i<n[1];i++){
        int u,v; scanf("%d%d",&u,&v);
        g[1][u].push_back(v); g[1][v].push_back(u);
    }
}

int ord[260];

void NTT(int *d,int len,int dr){
    for(int i=0;i<len;i++) if(ord[i] < i) swap(d[i],d[ord[i]]);
    for(int i=1;i<len;i<<=1){
        int wn = fast_pow(gg,(mod-1)/(2*i));
        if(dr == -1) wn = fast_pow(wn,mod-2);
        for(int j=0;j<len;j+=(i<<1)){
            for(int k=0,w=1;k<i;k++,w = (1ll*w*wn)%mod){
                int x = d[j+k],y = (1ll*w*d[j+k+i])%mod;
                d[j+k] = (x+y)%mod;
                d[j+k+i] = (x-y+mod)%mod;
            }
        }
    }
    if(dr == -1){
        int iv = fast_pow(len,mod-2);
        for(int i=0;i<len;i++) d[i] = (1ll*d[i]*iv)%mod;
    }
}

int A[260],B[260];
int fi[260],A0[260];

void INV(){
    int len = 1,bit = 0; while(len <= k) len<<=1,bit++;
    memset(A0,0,sizeof(A0));memset(fi,0,sizeof(fi));
    A0[0] = 1;
    for(int i=2,j=1;i<=len;i<<=1,j++){
        for(int k=0;k<i;k++) fi[k] = A[k];
        int rl = i*2,rb = j+1;
        for(int k=0;k<rl;k++) ord[k] = (ord[k>>1]>>1) + ((k&1)<<rb-1);
        NTT(A0,rl,1); NTT(fi,rl,1);
        for(int k=0;k<rl;k++){
            A0[k] = (2*A0[k]-(1ll*fi[k]*A0[k]%mod)*A0[k]%mod)%mod;
            if(A0[k] < 0) A0[k] += mod;
        }
        NTT(A0,rl,-1);
        for(int k=i;k<rl;k++) A0[k] = fi[k] = 0;
    }
    for(int i=0;i<=k;i++) A[i] = A0[i];
}

void dfs1(int kd,int now,int fa){
    for(auto it:g[kd][now]){
        if(it == fa) continue;
        dfs1(kd,it,now);
    }
    memset(A,0,sizeof(A));
    for(auto it:g[kd][now]){
        if(it == fa) continue;
        for(int i=0;i<=k-2;i+=2) A[i+2] = (A[i+2]+T[kd][it][i])%mod;
    }
    A[0] -= 1; if(A[0] < 0) A[0] += mod;
    for(int i=0;i<=k;i++){ A[i] *= -1; if(A[i] < 0) A[i] += mod;}
    INV();
    for(int i=0;i<=k;i++) T[kd][now][i] = A[i];
}

void dfs2(int kd,int now,int fa){
    memset(B,0,sizeof(B));
    for(auto it:g[kd][now]){
        if(it == fa) continue;
        for(int i=0;i<=k-2;i+=2) B[i+2] = (B[i+2]+T[kd][it][i])%mod;
    }
    for(int i=0;i<=k-2;i+=2) B[i+2] = (B[i+2]+gi[kd][now][i])%mod;
    for(auto it:g[kd][now]){
        if(it == fa) continue;
        for(int i=0;i<=k-2;i+=2) B[i+2] = (B[i+2]+mod-T[kd][it][i])%mod;
        memset(A,0,sizeof(A));
        for(int i=0;i<=k;i++) A[i] = (mod-B[i])%mod; A[0] = (1-A[0]+mod)%mod;
        INV(); for(int i=0;i<=k;i++) gi[kd][it][i] = A[i];
        for(int i=0;i<=k-2;i+=2) B[i+2] = (B[i+2]+T[kd][it][i])%mod;
    }
    memset(A,0,sizeof(A));
    for(int i=0;i<=k;i++) A[i] = (mod-B[i])%mod; A[0] = (1-A[0]+mod)%mod;
    INV(); for(int i=0;i<=k;i++) f[kd][now][i] = A[i];
    for(auto it:g[kd][now]){
        if(it == fa) continue;
        dfs2(kd,it,now);
    }
}

void solve(int kd){
    dfs1(kd,1,0);
    dfs2(kd,1,0);
}

void work(){
    solve(0);
    solve(1);
    for(int i=1;i<=k;i++){
        C[i][0] = C[i][i] = 1;
        for(int j=1;j<i;j++) C[i][j] = (C[i-1][j] + C[i-1][j-1]) % mod;
    }
    int ans = 0;
    for(int i=0;i<=k;i++){
        int s1 = 0,s2 = 0;
        for(int j=1;j<=n[0];j++) s1 += f[0][j][i],s1 %= mod;
        for(int j=1;j<=n[1];j++) s2 += f[1][j][k-i],s2 %= mod;
        int pp = (1ll*s1*s2)%mod;pp = (1ll*C[k][i]*pp)%mod;
        ans += pp; ans %= mod;
    }
    printf("%d",ans);
}

int main(){
    read();
    work();
    return 0;
}