Codeforces 1118F1 Tree Cutting (Easy Version) (DFS)

题面链接：传送门

F1. Tree Cutting (Easy Version)

time limit per test

2 seconds

memory limit per test

256 megabytes

input

standard input

output

standard output

You are given an undirected tree of nn 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?

Input

The first line contains a single integer nn (2≤n≤3⋅1052≤n≤3⋅105) — the number of vertices in the tree.

The second line contains nn integers a1,a2,…,ana1,a2,…,an (0≤ai≤20≤ai≤2) — the colors of the vertices. ai=1ai=1 means that vertex ii is colored red, ai=2ai=2 means that vertex ii is colored blue and ai=0ai=0 means that vertex ii is uncolored.

The ii-th of the next n−1n−1 lines contains two integers vivi and uiui (1≤vi,ui≤n1≤vi,ui≤n, vi≠uivi≠ui) — 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.

Output

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

Examples

input

5
2 0 0 1 2
1 2
2 3
2 4
2 5

output

1

input

5
1 0 0 0 2
1 2
2 3
3 4
4 5

output

4

input

3
1 1 2
2 3
1 3

output

0

题意描述：

       给你一颗树，树上结点可以染成红色和蓝色和无色，至少有一个红色和蓝色，如果一条边被切段，树变成两部分，如果每一部分只有一种颜色（无色不算颜色），那么，这条边是好的，现在问你有多少条好的边。

题目分析：

       因为只有两种颜色，所以输入的时候先计算每种颜色的个数，然后暴力dfs，可以得到其中一部分的染色情况，如果这部分只有一种颜色并且该颜色数等于总颜色数，那么ans++。具体看代码。

代码：

#include
using namespace std;
const int maxn = 300005;
int A[maxn], B[maxn][3], ans = 0, C[3], n;
vectorG[maxn];
void dfs(int u, int fa){
    B[u][A[u]]++;
    for(int v: G[u]){
        if(v != fa){
            dfs(v, u);
            B[u][1] += B[v][1];
            B[u][2] += B[v][2];
        }
    }
    if(B[u][1] && B[u][2]) return ;
    if(B[u][1] == C[1]|| B[u][2] == C[2]) ans++;
}
int main(){
    cin >> n;
    for(int i=1;i<=n;++i){
        scanf("%d", &A[i]);
        C[A[i]]++;
    }
    int u, v;
    for(int i=1;i