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);
 }