牛客多校day4-D.Rebuild Tree

Rebuild Tree

题意：

一个$n$个结点的树，删掉$k$条边，加上$k$条边，输出还是一个树的方案数模$998244353$。

思路：

删$k$条边之后形成$k+1$个连通块，设每个连通块的大小为$s_i$，则生成树的个数为$n^{k-1}\prod _{i=1}^{k+1}{s}_i$。
考虑如何求$\prod _{i=1}^{k+1}{s}_i$,用$d{p}_{i,j,0/1}$表示以$i$为根的子树删$j$条边，选没选结点的方案数。

代码：

#include<bits/stdc++.h>
#define fi first
#define se second
#define int long long
#define mp make_pair
#define pb push_back
#define ls x<<1
#define rs x<<1|1
#define lson x<<1,l,mid
#define rson x<<1|1,mid+1,r
#define pii pair<int,int>
#define all(x) x.begin(),x.end()
#define cl(x,y) memset(x,y,sizeof(x))
#define nxtp(a,n) next_permutation(a+1,a+n+1)
#define mem(x,y,n) memset(x,y,sizeof(int)*(n+5))
const int N=5e4+10;
const int mod=998244353;
const int inf=0x3f3f3f3f;
const double eps=1e-8;
const double pi=acos(-1);
const double INF=1e18;
using namespace std;
int n,k;
int dp[N][110][2],sz[N],tmp[110][2];
vector<int> e[N];
int qpow(int a,int b)
{
	int ans=1;
	while(b)
	{
		if(b&1)
			ans=ans*a%mod;
		a=a*a%mod;
		b>>=1;
	}
	return ans;
}
void dfs(int u,int fa)
{
	int i,j;
	sz[u]=1;
	dp[u][0][0]=dp[u][0][1]=1;
	for(auto v:e[u])
	{
		if(v==fa)
			continue;
		dfs(v,u);
		cl(tmp,0);
		for(i=0;i<=min(sz[u],k);i++)
			for(j=0;j<=min(sz[v],k);j++)
			{
				if(i+j>k)
					break;
				tmp[i+j][0]=(tmp[i+j][0]+dp[u][i][0]*dp[v][j][0]%mod)%mod;
				tmp[i+j][1]=(tmp[i+j][1]+dp[u][i][0]*dp[v][j][1]%mod)%mod;
				tmp[i+j][1]=(tmp[i+j][1]+dp[u][i][1]*dp[v][j][0]%mod)%mod;
				if(i+j+1>k)
					break;
				tmp[i+j+1][0] = (tmp[i+j+1][0] + dp[u][i][0] * dp[v][j][1] % mod )%mod;
                tmp[i+j+1][1] = (tmp[i+j+1][1] + dp[u][i][1] * dp[v][j][1] % mod )%mod;
			} 
		sz[u]+=sz[v];
		memcpy(dp[u],tmp,sizeof(dp[u]));
	}
}
signed main()
{
	ios::sync_with_stdio(false);
	cin.tie(0);cout.tie(0);
	int i;
	cin>>n>>k;
	for(i=1;i<n;i++)
	{
		int u,v;
		cin>>u>>v;
		e[u].pb(v);
		e[v].pb(u);
	}
	dfs(1,0);
	int ans=qpow(n,k-1)*dp[1][k][1]%mod;
	cout<<ans<<endl; 
	return 0;
}