蓝题训练9.28P15970 树上检查

P15970 树上检查
Tag： dp 思维
题意：略

题解：树上dp即可

参考代码

#include <bits/stdc++.h>
using namespace std;
using ll = long long;
const ll p= 998244353;
const ll inv2=499122177;
const ll inv3=332748118;
const ll inv4=748683265;
ll dp[200005][5],sz[200005];
void solve()
{
	ll k,n;
	cin>>n>>k;
	vector<vector<ll>> a(n+1);
	vector<vector<ll>> dp(n+1,vector<ll>(5,0)) ;
	vector<vector<ll>> dp2(n+1,vector<ll>(5,0)) ;
	vector<ll> sz(n+1);
	for(ll x,y,i=1;i<n;i++){
		cin>>x>>y;
		a[x].push_back(y);
		a[y].push_back(x);
	}
	auto dfs=[&](this auto &&dfs,ll x,ll fa) ->void{
		sz[x]=1;
		for(auto j:a[x]){
			if(j==fa) continue;
			dfs(j,x);
			sz[x]+=sz[j];
			dp2[x][2]=(dp2[j][2]+dp2[x][2])%p;
		}
		dp[x][1]=dp2[x][1]=sz[x];
		dp[x][2]=sz[x]*(sz[x]-1)/2%p;
		dp2[x][2]=(dp2[x][2]+sz[x]-1+p)%p;
	};
	auto dfs1=[&](this auto &&dfs1,ll x,ll fa) ->void{
		ll sum1=0,sum2=0,sum3=0,sum4=0;
		for(auto j:a[x]){
			if(j==fa) continue;
			dfs1(j,x);
			dp[x][3]=(dp[x][3]+dp[j][3])%p;
			dp[x][3]=(dp[x][3]+dp2[j][2])%p;
			sum1=(sum1+dp2[j][2])%p;
			sum2=(sum2+dp2[j][2]*dp[j][1]%p)%p;
			sum3=(sum3+sz[j])%p;
			sum4=(sum4+sz[j]*sz[j]%p)%p;
		}
		dp[x][3]=(dp[x][3]+(sz[x]-1)*sum1%p)%p;
		dp[x][3]=((dp[x][3]-sum2)%p+p)%p;
		dp[x][3]=(dp[x][3]+((sum3*sum3%p-sum4)%p+p)%p*inv2%p)%p;
	};
	auto dfsp=[&](this auto &&dfsp,ll x,ll fa) ->void{
		ll sum1=0,sum2=0,sum3=0,sum4=0;
		for(auto j:a[x]){
			if(j==fa) continue;
			dfsp(j,x);
			dp2[x][3]=(dp2[x][3]+dp2[j][3])%p;
			dp2[x][3]=(dp2[x][3]+dp2[j][2])%p;
		}
	};
	auto dfs2=[&](this auto &&dfs2,ll x,ll fa) ->void{
		ll sum1=0,sum2=0,sum3=0,sum4=0,sum6=0,sum5=0;
		for(auto j:a[x]){
			if(j==fa) continue;
			dfs2(j,x);
			dp[x][4]=(dp[x][4]+dp[j][4])%p;
			dp[x][4]=(dp[x][4]+dp2[j][3])%p;
			sum1=(sum1+dp2[j][3])%p;
			sum2=(sum2+dp2[j][3]*dp[j][1]%p)%p;
			sum3=(sum3+sz[j])%p;
			sum4=(sum4+sz[j]*dp2[j][2]%p)%p;
			sum5=(sum5+dp2[j][2])%p;
			sum6=(sum6+dp2[j][2]*dp2[j][2]%p)%p;
		}
		dp[x][4]=(dp[x][4]+(sz[x]-1)*sum1%p)%p;
		dp[x][4]=((dp[x][4]-sum2)%p+p)%p;
		dp[x][4]=(dp[x][4]+((sum5*sum5%p-sum6)%p+p)%p*inv2%p)%p;
		dp[x][4]=(dp[x][4]+((sum3*sum5%p-sum4)%p+p)%p)%p;
	};
	if(k==1) cout<<n<<"\n";
	else if(k==2) cout<<n*(n-1)/2%p<<"\n";
	else if(k==3){
		dfs(1,0);
		dfs1(1,0);
		ll ans=n*(n-1)%p*(n-2)%p*inv2%p*inv3%p;
		cout<<(dp[1][3]+(ans-dp[1][3]+p)%p*2%p)%p<<"\n";
	}
	else if(k==4){
		dfs(1,0);
		dfsp(1,0);
		dfs2(1,0);
		ll ans=n*(n-1)%p*(n-2)%p*(n-3)%p*inv2%p*inv3%p*inv4%p;
		cout<<(dp[1][4]+(ans-dp[1][4]+p)%p*2%p)%p<<"\n";
	}
}
int main()
{
	ios::sync_with_stdio(0);
	cin.tie(0);
	ll T = 1;
	cin >> T;
	while (T--)
		solve();
	return 0;
}

---