与我前面写过的 战略游戏这题 求摆放方案很像
树形DP
f[u][c]
: 表示以 $u$ 为根,$u$ 涂上颜色为 $c$ 的方案数
- 如果 $u$ 的颜色为 $3$, 那么子树只能涂 $1$, $2$, 那么子树 $j$ 的方案就是
f[j][1] 和 f[j][2]
- 那么对于所有子树而言, 拼在一起的方案就是
:
(f[j][1] * f[j + 1][1] * f[j + 1][2] ... + f[j][2] * f[j + 1][1] * f[j + 1][2] ... )
- 即为
(f[j][1] + f[j][2]) * (f[j + 1][1] + f[j + 1][2]) * ... )
#include <cstring>
#include <iostream>
#include <algorithm>
using namespace std;
const int N = 1e5 + 10, M = N + N, MOD = 1e9 + 7;
int n, m;
int c[N];
int h[N], e[M], ne[M], idx;
int 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)
{
f[u][1] = f[u][2] = f[u][3] = 1;
for (int i = 1; i <= 3; i ++ ) // 先将不合法状态去掉
if (c[u] && c[u] != i) f[u][i] = 0; // 如果 u 本来有颜色,那么涂抹其它颜色是不合法的
for (int i = h[u]; ~i; i = ne[i])
{
int j = e[i];
if (j == fa) continue;
dfs(j, u);
f[u][1] = (1ll * f[u][1] * (f[j][2] + f[j][3])) % MOD; // 注意会爆long long
f[u][2] = (1ll * f[u][2] * (f[j][1] + f[j][3])) % MOD; // 不合法状态是0, 不会更新状态
f[u][3] = (1ll * f[u][3] * (f[j][1] + f[j][2])) % MOD;
}
}
int main()
{
scanf("%d%d", &n, &m);
memset(h, -1, sizeof h);
for (int i = 0; i < n - 1; i ++ )
{
int a, b;
scanf("%d%d", &a, &b);
add(a, b), add(b, a);
}
while (m -- )
{
int a, b;
scanf("%d%d", &a, &b);
c[a] = b;
}
dfs(1, -1);
printf("%d\n", (f[1][1] + f[1][2] + f[1][3]) % MOD);
return 0;
}