题库练习 Tree Cutting (Easy Version)
← 上一题 下一题 →

A12431 | Tree Cutting (Easy Version)

时间限制1s
内存限制256MB
通过 / 提交0/0

题目描述

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?

输入格式

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.

输出格式

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

输入输出样例

输入 #1
5
2 0 0 1 2
1 2
2 3
2 4
2 5
输出 #1
1
输入 #2
5
1 0 0 0 2
1 2
2 3
3 4
4 5
输出 #2
4
输入 #3
3
1 1 2
2 3
1 3
输出 #3
0
C++ 编辑器
输入
输出