CF1118F1.Tree Cutting (Easy Version)

You are given an undirected tree of $n$ vertices.

Some vertices are colored blue, some are colored red and some are uncolored. It is guaranteed that the tree contains at least one red vertex and at least one blue vertex.

You choose an edge and remove it from the tree. Tree falls apart into two connected components. Let's call an edge nice if neither of the resulting components contain vertices of both red and blue colors.

How many nice edges are there in the given tree?

给你一棵包含 $n$ 个顶点的无向树。

其中一些顶点被染成蓝色，一些被染成红色，还有一些未染色。保证该树中至少存在一个红色顶点和一个蓝色顶点。

你选择一条边并将其从树中删除，树将分裂为两个连通分量。若删除某条边后得到的两个连通分量中，均不同时包含红色和蓝色顶点，则称该边为“好边”（nice edge）。

问：给定树中一共有多少条好边？

The first line contains a single integer $n$ ($2 \le n \le 3 \cdot 10^5$) — the number of vertices in the tree.

The second line contains $n$ integers $a_1, a_2, \dots, a_n$ ($0 \le a_i \le 2$) — the colors of the vertices. $a_i = 1$ means that vertex $i$ is colored red, $a_i = 2$ means that vertex $i$ is colored blue and $a_i = 0$ means that vertex $i$ is uncolored.

The $i$-th of the next $n - 1$ lines contains two integers $v_i$ and $u_i$ ($1 \le v_i, u_i \le n$, $v_i \ne u_i$) — the edges of the tree. It is guaranteed that the given edges form a tree. It is guaranteed that the tree contains at least one red vertex and at least one blue vertex.

第一行包含一个整数 $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$ 行中，第 $i$ 行包含两个整数 $v_i$ 和 $u_i$（$1 \le v_i, u_i \le n$，且 $v_i \ne u_i$）——树的边。保证所给边构成一棵树。同时保证该树中至少包含一个红色顶点和一个蓝色顶点。

Print a single integer — the number of nice edges in the given tree.

输出一个整数——给定树中“好边”的数量。

Here is the tree from the first example:

The only nice edge is edge $(2, 4)$. Removing it makes the tree fall apart into components ${4}$ and ${1, 2, 3, 5}$. The first component only includes a red vertex and the second component includes blue vertices and uncolored vertices.

Here is the tree from the second example:

Every edge is nice in it.

Here is the tree from the third example:

Edge $(1, 3)$ splits the into components ${1}$ and ${3, 2}$, the latter one includes both red and blue vertex, thus the edge isn't nice. Edge $(2, 3)$ splits the into components ${1, 3}$ and ${2}$, the former one includes both red and blue vertex, thus the edge also isn't nice. So the answer is 0.

第一个样例对应的树如下：

其中唯一的“优美边”是边 $(2, 4)$。删除该边后，树被分割为两个连通分量：$\{4\}$ 和 $\{1, 2, 3, 5\}$。第一个分量仅包含一个红色顶点，第二个分量则同时包含蓝色顶点和未着色顶点。

第二个样例对应的树如下：

该树中的每条边都是“优美边”。

第三个样例对应的树如下：

边 $(1, 3)$ 将树分割为分量 $\{1\}$ 和 $\{3, 2\}$，后者同时包含红色顶点和蓝色顶点，因此该边不是“优美边”。边 $(2, 3)$ 将树分割为分量 $\{1, 3\}$ 和 $\{2\}$，前者同时包含红色顶点和蓝色顶点，因此该边也不是“优美边”。故答案为 $0$。