思路
为了方便,把不是红色节点的节点认为是黑色节点。考虑只有一棵子树,如果父节点是黑色,则它必须和子节点分到同一个连通块,且该连通块包含且只包含一个红色节点。那么对子树的分割要求就是所有连通块包含且只包含一个红色节点。如果父节点是红色,如果让它单独作为一个连通块,那么对子树的分割要求就是所有连通块包含且只包含一个红色节点。如果让父节点和子节点处于相同连通块,那么对子树的分割要求就是子节点所处连通块必须全黑,其它连通块包含且只包含一个红色节点。除此之外,没有其它可能性。当考虑有多棵子树时,可以分步进行,先将父节点和一棵子树合并,之后把父节点和它所在连通块(或其单独作为一个连通块)作为整体考虑(当成一个节点)。若该连通块全黑,就认为它是一个黑色节点,否则认为它是一个红色节点,接着用这个节点去合并下一棵子树,这个过程是递归的。
用 $f(u, 1)$ 表示将以 $u$ 为根的树,切成若干连通块,每个连通块包含且只包含 $1$ 个红色节点(可以包含多个黑色节点)的切法数量。用 $f(u, 0)$ 表示将以 $u$ 为根的树,切成若干连通块,其中 $u$ 所在连通块只包含黑色节点,其它连通块(这部分可以不存在)包含且只包含 $1$ 个红色节点(可以包含多个黑色节点)的切法数量。答案为 $f(root, 1)$。
假设当前已经考虑了 $u$ 的一些子节点,并将结果组合到了 $f(u)$,现在考虑子节点 $v$:
- 考虑 $v$ 对 $f(u, 0)$ 的贡献。首先,从此刻 $f(u, 0)$ 对应的集合拿出任意一种切法,$u$ 所在连通块只包含黑色节点。(对于 $f(u, 1)$ 中的任意切法,$u$ 已经与 $1$ 个红色点连通,和 $f(v)$ 组合不可能产生满足 $f(u, 0)$ 的方案)。
- 若从 $f(v, 0)$ 对应集合拿出任意一种切法,$v$ 所在连通块只包含黑色节点。让 $u, v$ 连通,则 $u$ 所在连通块只包含黑色节点,而其它连通块包含且只包含 $1$ 个红色节点,因此可以使 $f(u, 0)$ 加一,而这样的组合有 $f(u, 0) \times f(v, 0)$ 种。若不让 $u, v$ 连通,则存在 $2$ 个全黑连通块,不满足 $f(u, 0)$ 要求。因此
$$ f(u, 0) \leftarrow f(u, 0) \times f(v, 0) $$ - 若从 $f(v, 1)$ 对应集合拿出任意一种切法,$v$ 所在连通块包含 $1$ 个红色节点。要让组合后的 $u$ 所在连通块不包含红色节点,就不能让 $u, v$ 连通。而不连通时,两个集合任意组合都满足 $f(u, 0)$ 的要求,因此
$$ f(u, 0) \leftarrow f(u, 0) \times f(v, 1) $$
- 若从 $f(v, 0)$ 对应集合拿出任意一种切法,$v$ 所在连通块只包含黑色节点。让 $u, v$ 连通,则 $u$ 所在连通块只包含黑色节点,而其它连通块包含且只包含 $1$ 个红色节点,因此可以使 $f(u, 0)$ 加一,而这样的组合有 $f(u, 0) \times f(v, 0)$ 种。若不让 $u, v$ 连通,则存在 $2$ 个全黑连通块,不满足 $f(u, 0)$ 要求。因此
- 考虑 $v$ 对 $f(u, 1)$ 的贡献。
- 首先,从 $f(u, 0)$ 对应集合拿出任意一种切法,$u$ 所在连通块只包含黑色节点。从 $f(v, 1)$ 对应集合拿出任意一种切法,$v$ 所在连通块包含 $1$ 个红色节点。让 $u, v$ 连通,则每个连通块都包含且只包含 $1$ 个红色节点,因此可以使 $f(u, 1)$ 加一,而这样的组合有 $f(u, 0) \times f(v, 1)$ 种。若不让 $u, v$ 连通,则 $u$ 所在连通块无法包含红色节点。此外,$f(v, 0)$ 所在集合与 $f(u, 0)$ 组合无法对 $f(u, 1)$ 做出贡献。因此
$$ f(u, 1) \leftarrow f(u, 0) \times f(v, 1) $$ - 从 $f(u, 1)$ 对应集合拿出任意一种切法,$u$ 所在连通块包含 $1$ 个红色节点。
- 若从 $f(v, 0)$ 对应集合拿出任意一种切法,$v$ 所在连通块只包含黑色节点。让 $u, v$ 连通,则所有连通块都只包含 $1$ 个红色节点,可以使 $f(u, 1)$ 加一,而这样的组合有 $f(u, 1) \times f(v, 0)$ 种。若不让 $u, v$ 连通,则存在全黑连通块,不满足 $f(u, 1)$ 要求。因此
$$ f(u, 1) \leftarrow f(u, 1) \times f(v, 0) $$ - 若从 $f(v, 1)$ 对应集合拿出任意一种切法,一定不能让 $u, v$ 连通,且若让它们不连通,则所有连通块都只包含 $1$ 个红色节点,而这两个集合可以任意组合,因此
$$ f(u, 1) \leftarrow f(u, 1) \times f(v, 1) $$
- 若从 $f(v, 0)$ 对应集合拿出任意一种切法,$v$ 所在连通块只包含黑色节点。让 $u, v$ 连通,则所有连通块都只包含 $1$ 个红色节点,可以使 $f(u, 1)$ 加一,而这样的组合有 $f(u, 1) \times f(v, 0)$ 种。若不让 $u, v$ 连通,则存在全黑连通块,不满足 $f(u, 1)$ 要求。因此
- 首先,从 $f(u, 0)$ 对应集合拿出任意一种切法,$u$ 所在连通块只包含黑色节点。从 $f(v, 1)$ 对应集合拿出任意一种切法,$v$ 所在连通块包含 $1$ 个红色节点。让 $u, v$ 连通,则每个连通块都包含且只包含 $1$ 个红色节点,因此可以使 $f(u, 1)$ 加一,而这样的组合有 $f(u, 0) \times f(v, 1)$ 种。若不让 $u, v$ 连通,则 $u$ 所在连通块无法包含红色节点。此外,$f(v, 0)$ 所在集合与 $f(u, 0)$ 组合无法对 $f(u, 1)$ 做出贡献。因此
若 $u$ 为叶节点,根据定义 $f(u, color[u]) = 1, f(u, !color[u]) = 0$。注意:$f(u, 1)$ 依赖的 $f(u, 0)$ 是没有被 $f(v)$ 更新过的,因此要先算 $f(u, 1)$。
#include <cstring>
#include <iostream>
using namespace std;
using LL = long long;
const int N = 100010, M = N << 1, MOD = 1e9 + 7;
int n;
int c[N];
int e[M], ne[M], h[N], idx;
int f[N][2];
void add(int a, int b) { e[idx] = b, ne[idx] = h[a], h[a] = idx++; }
void dfs(int u, int p) {
f[u][c[u]] = 1;
for (int i = h[u]; ~i; i = ne[i]) {
int v = e[i];
if (v == p) continue;
dfs(v, u);
LL t = f[v][0] + f[v][1];
f[u][1] = ((LL)f[u][0] * f[v][1] % MOD + f[u][1] * t % MOD) % MOD;
f[u][0] = f[u][0] * t % MOD;
}
}
int main() {
scanf("%d", &n);
for (int i = 1; i <= n; i++) scanf("%d", c + i);
memset(h, -1, sizeof h);
for (int i = 1, a, b; i < n; i++) {
scanf("%d%d", &a, &b);
add(a, b), add(b, a);
}
dfs(1, 0);
printf("%d\n", f[1][1]);
return 0;
}