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