2021牛客多校#4 D-Rebuild Tree(Prufer序列)(树形DP)

$\mathcal{D}escription$

一棵树删k条边，再加k条边还是树，求方案数

$\mathcal{S}olution$

前置知识：prufer序列

• 设有k个连通块，k-1条边，则prufer序列长度为k-2
• 每个位置可填1~n$\Rightarrow n^{k-2}$
• 每条边联通连通块时可以连$s_i$个节点中任意一个
• 综上可得

$n^{k-2}\ast \prod _{i=1}^k{s}_i$ (k为联通块数量，$s_i$为联通块大小)

• 本题中$n^{k-1}\ast \prod _{i=1}^{k+1}{s}_i$(k条边k+1个联通块)

求$\prod _{i=1}^k{s}_i$：树形DP

问题等价于$n$点树分成$k+1$个连通块，并在每个连通块中任选一点的方案数
设$d{p}_{i,j,t}$表示以i为根的子树中分为j个连通块,i点所在块中是否已选点方案数

$$\{\begin{array}{rl}t& =0\Rightarrow 未选点\\ t& =1\Rightarrow 已选点\end{array}$$

for son
	for j i的联通块数
		for k son的连通块数

$task:将son子树中连通块合并到i上$

$Case1：i块中的点在i的祖先处已经选好\Rightarrow d{p}_{i,j,0}$

• $i与son不删边\Rightarrow son中不选点，并入i中$
  $⟹d{p}_{i,j,0}+=d{p}_{i,j-k+1,0}\ast d{p}_{son,k,0}$

• $i与son删边，不管son\Rightarrow son中必须选点（每块至少一点），并入i中$
  $⟹d{p}_{i,j,0}+=d{p}_{i,j-k,0}\ast d{p}_{son,k,1}$

$Case2：i块中的点在i的子树中\Rightarrow d{p}_{i,j,1}$

• $不删边，点在i这块中$
  $⟹d{p}_{i,j,1}+=d{p}_{i,j-k+1,1}\ast d{p}_{son,k,0}$

• $不删边，点在son这块中$
  $⟹d{p}_{i,j,1}+=d{p}_{i,j-k+1,0}\ast d{p}_{son,k,1}$

• $删边隔离，i和son中都有点$
  $⟹d{p}_{i,j,1}+=d{p}_{i,j-k,1}\ast d{p}_{son,k,1}$

#include <bits/stdc++.h>
#define ll long long
using namespace std;
const int N=5e4+10,K=110,M=998244353;
int n,k,sz[N],tot,h[N];
struct Edge{int v,nxt;}e[N<<1];
void add(int a,int b){e[++tot].v=b,e[tot].nxt=h[a],h[a]=tot;}
int dp[N][K][2],pre[K][2];
void dfs(int x,int fx){
	sz[x]=1;
	dp[x][1][0]=dp[x][1][1]=1;
	for(int i=h[x],v;i;i=e[i].nxt)if((v=e[i].v)!=fx)
	{
		dfs(v,x);
		memcpy(pre,dp[x],sizeof(pre));
		memset(dp[x],0,sizeof(dp[x]));
		for(int y=1;y<=min(sz[x],k+1);y++)
			for(int z=1;z<=min(sz[v],k+1)&&y+z-1<=k+1;z++)
				dp[x][y+z-1][0]=(dp[x][y+z-1][0]+1ll*pre[y][0]*dp[v][z][0])%M,
				dp[x][y+z][0]=(dp[x][y+z][0]+1ll*pre[y][0]*dp[v][z][1])%M,
				dp[x][y+z-1][1]=(dp[x][y+z-1][1]+1ll*pre[y][1]*dp[v][z][0])%M,
				dp[x][y+z-1][1]=(dp[x][y+z-1][1]+1ll*pre[y][0]*dp[v][z][1])%M,
				dp[x][y+z][1]=(dp[x][y+z][1]+1ll*pre[y][1]*dp[v][z][1])%M;
		sz[x]+=sz[v];
	}
}
int main(){
	scanf("%d%d",&n,&k);
	for(int i=1,x,y;i<n;i++) 
		scanf("%d%d",&x,&y),add(x,y),add(y,x);
	dfs(1,0);
	int ans=dp[1][k+1][1];
	for(int i=1;i<k;i++)
		ans=1ll*ans*n%M;
	printf("%d\n",ans);
}