CF1118F1 Tree Cutting (Easy Version)

给定一棵包含 $n$ 个顶点的无向树。 有些顶点被染成蓝色，有些被染成红色，还有一些未染色。保证树中至少有一个红色顶点和至少一个蓝色顶点。 你可以选择一条边并将其从树中移除。此时树会分裂成两个连通分量。如果分裂后得到的两个连通分量中都不同时包含红色和蓝色顶点，则称这条边是“好边”。 请问在给定的树中，有多少条好边？

第一行包含一个整数 $n$（$2 \le n \le 3 \cdot 10^5$），表示树的顶点数。 第二行包含 $n$ 个整数 $a_1, a_2, \dots, a_n$（$0 \le a_i \le 2$），表示每个顶点的颜色。$a_i = 1$ 表示第 $i$ 个顶点为红色，$a_i = 2$ 表示第 $i$ 个顶点为蓝色，$a_i = 0$ 表示第 $i$ 个顶点未染色。 接下来 $n-1$ 行，每行包含两个整数 $v_i$ 和 $u_i$（$1 \le v_i, u_i \le n$，$v_i \ne u_i$），表示树中的一条边。保证给定的边构成一棵树，且树中至少有一个红色顶点和至少有一个蓝色顶点。

输出一个整数，表示树中好边的数量。

以下是第一个样例中的树：  唯一的好边是边 $(2, 4)$。移除该边后，树分裂为 $\{4\}$ 和 $\{1, 2, 3, 5\}$ 两个连通分量。第一个分量只包含红色顶点，第二个分量包含蓝色顶点和未染色顶点。 以下是第二个样例中的树：  其中每一条边都是好边。 以下是第三个样例中的树：  边 $(1, 3)$ 将树分成 $\{1\}$ 和 $\{3, 2\}$ 两个分量，后者同时包含红色和蓝色顶点，因此该边不是好边。边 $(2, 3)$ 将树分成 $\{1, 3\}$ 和 $\{2\}$，前者同时包含红色和蓝色顶点，因此该边也不是好边。所以答案为 $0$。 由 ChatGPT 4.1 翻译