蓝魔法师(树形DP)

蓝魔法师(树形DP)

思路：以每条边进行状态转移，令 $dp[i][j]$以 $i$为根的子树 $i$所在连通块大小为 $j$的方案数。

对于边 $edge(u,v)$，若不隔断该边，则 $dp[u][i+j]=\sum dp[u][i]\times dp[v][j]\ (i+j\le k)$

若隔断该边，则 $dp[u][i]=dp[u][i]\times \sum\limits_{j=1}^{k}dp[v][j],i\in[1,size[u]]$

为了方便储存我们可以用 $dp[v][0]=\sum\limits_{j=1}^{k}dp[v][j]$

最后答案为： $\sum\limits_{i=1}^{k}dp[1][i]=dp[1][0]$

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
const int N=2e3+5,M=2e4+5,inf=0x3f3f3f3f,mod=998244353;
#define mst(a,b) memset(a,b,sizeof a)
#define lx x<<1
#define rx x<<1|1
#define reg register
#define PII pair<int,int>
#define fi first
#define se second
#define pb push_back
int n,k,sz[N];
ll dp[N][N],tmp[N];
vector<int>e[N];
void dfs(int u,int fa){
    sz[u]=dp[u][1]=1;
    for(int v:e[u]){
        if(v==fa) continue;
        dfs(v,u);
        for(int i=1;i<=sz[u];i++)
           for(int j=0;j<=min(sz[v],k-i);j++)
           tmp[i+j]=(tmp[i+j]+dp[u][i]*dp[v][j]%mod)%mod;
        for(int i=1;i<=k;i++)
           dp[u][i]=tmp[i],tmp[i]=0;
        sz[u]+=sz[v];
    }
    for(int i=1;i<=k;i++) dp[u][0]=(dp[u][0]+dp[u][i])%mod;
}
int main(){
    scanf("%d%d",&n,&k);
    for(int i=1,u,v;i<n;i++){
        scanf("%d%d",&u,&v),e[u].pb(v),e[v].pb(u);
    }
    dfs(1,0);
    printf("%lld\n",dp[1][0]);
	return 0;
}