Manthan, Codefest 17 C. Helga Hufflepuff's Cup（树形计数DP）（*）

题目链接


题意：给你一棵树，可以染 m 种颜色，现定义一种特殊的颜色 K ，一棵树上最多能有 X 个特殊颜色。如果一个节点为特殊颜色 K ，那么他相邻的节点的值只能选比 K 小的颜色，问一共有多少种染色方案。
思路：dp[i][j][k] 代表以 i 为根的子树中共选择了 j 个特殊颜色，且当前节点 i 的状态为 k 的染色方案数。

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
const int maxn=1e5+1; 
const int mod=1e9+7;
int n,m,u,v,K,X,size[maxn];
ll ans=0,dp[maxn][12][3],temp[12][3];
vector<int>g[maxn];
void dfs(int x,int fa)
{
	dp[x][0][0]=K-1;
	dp[x][1][1]=1;
	dp[x][0][2]=m-K;
	size[x]=1;
	for(int to:g[x])
	{
		if(to==fa) continue;
		dfs(to,x);
		memset(temp,0,sizeof(temp));
		for(int j=0;j<=size[x];++j)
		for(int k=0;k<=size[to];++k)
		if(j+k<=X)
		{
			temp[j+k][0] = (temp[j+k][0] + dp[x][j][0]*(dp[to][k][0] + dp[to][k][1] + dp[to][k][2])%mod)%mod;
                temp[j+k][1] = (temp[j+k][1] + dp[x][j][1]*(dp[to][k][0])%mod)%mod;
                temp[j+k][2] = (temp[j+k][2] + dp[x][j][2]*(dp[to][k][0] + dp[to][k][2])%mod)%mod;
		}
	   size[x] = min(size[x]+size[to],X);
        for(int j=0;j<=size[x];j++)
            for(int k=0;k<3;k++)
                dp[x][j][k] = temp[j][k],temp[j][k]=0;
	}
}
int main()
{
	scanf("%d %d",&n,&m);
	for(int i=1;i<n;++i)
	{
		scanf("%d %d",&u,&v);
		g[u].push_back(v);
		g[v].push_back(u);
	}
	scanf("%d %d",&K,&X);
	dfs(1,-1);
	for(int i=0;i<=X;++i)
	for(int j=0;j<3;++j)
	ans=(ans+dp[1][i][j])%mod;
	printf("%lld\n",ans);
 }