CodeForces - 855C Helga Hufflepuff's Cup 树形DP

传送门：Codeforces 855C

题意：给你一棵树，可以染m种颜色（1~m），现在定义一种特殊颜色k，一棵树上最多能有x个特殊颜色，如果一个节点为特殊颜色k，那么他相邻的节点的值只能选比k小的颜色。现在问你一共有多少种染色的方法。

思路：dp[i][j][k] 表示以i为根的子树上有j个特殊颜色结点的染色方案数，

k = 0 ：i结点的颜色比k小

k = 1 ：i结点的颜色为k

k = 2 ：i结点的颜色大于k

代码：

#include<bits/stdc++.h>
#define ll long long
#define inf 0x3f3f3f3f
using namespace std;
typedef pair<int,int> P;
const int MAXN = 100010;
const int mod = 1e9 + 7;
vector<int> mp[MAXN];
ll N, M, K, X;
ll dp[MAXN][11][3], sz[MAXN], tmp[11][3];
void dfs(int u, int fa)
{
	dp[u][0][0] = K - 1;
	dp[u][1][1] = 1;
	dp[u][0][2] = M - K;
	sz[u] = 1;
	int v;
	for(int i = 0; i < mp[u].size(); i++)
	{
		v = mp[u][i];
		if(v == fa) continue;
		dfs(v, u);
		memset(tmp, 0, sizeof(tmp));
		for(int j = 0; j <= sz[u]; j++)
		for(int k = 0; k <= sz[v]; k++)
		{
			if(j + k > X) continue;
			tmp[j + k][0] = (tmp[j + k][0] + dp[u][j][0] * (dp[v][k][0] + dp[v][k][1] + dp[v][k][2])) % mod;
			tmp[j + k][1] = (tmp[j + k][1] + dp[u][j][1] * dp[v][k][0]) % mod;
			tmp[j + k][2] = (tmp[j + k][2] + dp[u][j][2] * (dp[v][k][2] + dp[v][k][0])) % mod;
		}
		sz[u] = min(sz[u] + sz[v], X);
		for(int j = 0; j <= sz[u]; j++)
		for(int k = 0; k < 3; k++)
		dp[u][j][k] = tmp[j][k];
	}
}
int main()
{
	int u, v;
	ll ans = 0;
	cin >> N >> M;
	for(int i = 1; i < N; i++)
	{
		scanf("%d %d", &u, &v);
		mp[u].push_back(v), mp[v].push_back(u);
	}
	cin >> K >> X;
	dfs(1, -1);
	for(int i = 0; i <= sz[1]; i++)
	for(int j = 0; j < 3; j++)
	ans = (ans + dp[1][i][j]) % mod;
	cout << ans << endl;
 	return 0;
}