Codeforces Round #540 (Div. 3) F1. Tree Cutting (Easy Version) /*遍历树样题*/

 http://codeforces.com/contest/1118/problem/F1

Codeforces Round #540 (Div. 3) F1. Tree Cutting (Easy Version) /*遍历树样题*/

 

统计出以i为结点的子树的红色结点个数，和蓝色结点个数，判断一下就好了

（注意n==2，因为遍历出发找的是>=2度的点）

#include<bits/stdc++.h>
using namespace std;
#define LL long long
#define mod 1000000007
#define N 300005
/***************************************************************************************************************************************************************************************************/
int n,red,bule,a[N];
vector<int> v[N];
int dp[N][3];

int dfs(int root,int per,int c){
	if(a[root]==c)  dp[root][c]++;
	if(v[root].size()==1) return dp[root][c];
	for(int i=0;i<v[root].size();i++){
		if(v[root][i]==per) continue;
		dp[root][c]+=dfs(v[root][i],root,c);
	}
	return dp[root][c];
}

int main(){
	cin>>n;

	for(int i=1;i<=n;i++){
		scanf("%d",&a[i]);
		if(a[i]==1) red++;
		if(a[i]==2) bule++;
	}
	for(int i=1;i<=n-1;i++){
		int vv,uu;
		scanf("%d%d",&vv,&uu);
		v[vv].push_back(uu);
		v[uu].push_back(vv);
	}
	if(n==2) {
        cout<<1;
        return 0;
	}
	for(int i=1;i<=n;i++){
        if(v[i].size()>1){
            dfs(i,-1,1);
            dfs(i,-1,2);
            break;
        }
	}
	int ans=0;
	for(int i=1;i<=n;i++){
		if(dp[i][1]==red&&dp[i][2]==0) ans++;
		if(dp[i][1]==0&&dp[i][2]==bule) ans++;
	}
	cout<<ans;
    return 0;
}