url: https://www.luogu.com.cn/problem/P4084
tag:
动态规划，树形DP，USACO，2017
思路：
用树形DP来解决。定义数组 f[i][k] 表示以i为根节点且i涂的颜色为k时的方案数。然后根据k的不同分别更新即可。最后答案输出以1为根节点时3个颜色方案数之和（以其他节点为根节点也可以）。
代码：
#include <iostream>
#include <cstdio>
#include <cstring>
#include <algorithm>
using namespace std;
typedef long long LL;
const int N = 1e5 + 10;
int mod = 1e9 + 7;
int e[N * 2], ne[N * 2], h[N], idx;
int n, k;
int c[N];
LL f[N][4];
void add(int a, int b)
{
e[idx] = b;
ne[idx] = h[a];
h[a]= idx ++;
}
void dfs(int u, int fa)
{
if (c[u]) f[u][c[u]] = 1;
else
{
f[u][1] = f[u][2] = f[u][3] = 1;
}
for (int i = h[u]; ~i; i = ne[i])
{
int j = e[i];
if (j == fa) continue;
dfs(j, u);
f[u][1] = f[u][1] * ((f[j][2] + f[j][3]) % mod) % mod;
f[u][2] = f[u][2] * ((f[j][1] + f[j][3]) % mod) % mod;
f[u][3] = f[u][3] * ((f[j][1] + f[j][2]) % mod) % mod;
}
}
int main()
{
memset(h, -1, sizeof h);
cin >> n >> k;
for (int i = 0; i < n - 1; i ++)
{
int a, b;
cin >> a >> b;
add(a, b), add(b, a);
}
for (int i = 0; i < k; i++)
{
int a, b;
cin >> a >> b;
c[a] = b;
}
dfs(1, 0);
LL res = 0;
for (int i = 1; i <= 3; i ++) res = (res + f[1][i]) % mod;
cout << res << endl;
return 0;
}