[codeforces1118F1]Tree Cutting (Easy Version)

time limit per test : 2 seconds
memory limit per test : 256 megabytes

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?

Input

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. $ai=1$ means that vertex $i$ is colored red, $a_i=2$ means that vertex i is colored blue and $ai=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̸=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.

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

题意：
给一个树，树的节点有三种状态（$0$ : 没涂色,$1$ : 涂成红色,$2$ : 涂成蓝色），询问树中有多少条边满足：这棵树删除这条边之后分裂成两个联通块之后，这两个联通块都不是既有蓝色的节点也有红色节点的联通块。（数据保证至少有1个红色节点和一个蓝色节点）
题解：
我们可以知道，一条边是连接一个子树和树其他部分的桥，那么一条边是否是满足条件的边只需要满足这条边所对应的子树是否包含所有的红色点，或者包含所有的蓝色点。所以我们dfs的时候只需要返回一下子树的蓝色点个数和红色点个数就行了。然后交给上一级统计。

#include<bits/stdc++.h>
#define LiangJiaJun main
#define pa pair<int,int>
using namespace std;
int n,c[4],a[300004],ans;
int ne,h[300004];
bool vis[300004];
struct edge{
    int to,nt;
}e[600004];
void add(int u,int v){
     e[++ne].to=v;e[ne].nt=h[u];h[u]=ne;
}

pa  dfs(int x){
    vis[x]=1;
    pa g=make_pair(0,0);
    if(a[x]==1)g.first++;
    if(a[x]==2)g.second++;
    for(int i=h[x];i;i=e[i].nt){
        if(vis[e[i].to])continue;
        pa now=dfs(e[i].to);
        g.first+=now.first;
        g.second+=now.second;
        int fc[4];
        fc[1]=c[1]-now.first;
        fc[2]=c[2]-now.second;
        if(!((fc[1]>0&&fc[2]>0)||(now.first>0&&now.second>0)))++ans;
    }
    return g;
}
int w33ha(){
    ne=0;
    ans=0;
    memset(c,0,sizeof(c));
    memset(vis,0,sizeof(vis));
    memset(h,0,sizeof(h));
    for(int i=1;i<=n;i++){
        scanf("%d",&a[i]);
        ++c[a[i]];
    }
    for(int i=1;i<n;i++){
        int u,v;scanf("%d%d",&u,&v);
        add(u,v);
        add(v,u);
    }
    dfs(1);
    printf("%d\n",ans);
    return 0;
}
int LiangJiaJun(){
    while(scanf("%d",&n)!=EOF)w33ha();
    return 0;
}