Codeforces Round #540 (Div. 3)--1118F1

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

#include<bits/stdc++.h>
using namespace std;
int n;
vector<int> color;
vector<vector<int> > tree;
int red=0,blue=0;
int ans=0;

pair<int,int> dfs(int v,int p=-1){
    int r=(color[v]==1);
    int b=(color[v]==2);
    for(int i=0;i<tree[v].size();i++){
        int u=tree[v][i];
        if(u!=p){//避免回溯到上一个经过的点
            pair<int,int> tmp=dfs(u,v);
            ans+=(tmp.first==red&&tmp.second==0);//只有之前经过的点中包含了全部的红或者蓝点才可以
            ans+=(tmp.first==0&&tmp.second==blue);
            r+=tmp.first;//加上之前点中的红和蓝点
            b+=tmp.second;
        }
    }
    return make_pair(r,b);
}

int main(){
    cin>>n;
    color.resize(n);
    for(int i=0;i<n;i++){
        cin>>color[i];
        if(color[i]==1)
            red+=1;
        else if(color[i]==2)
            blue+=1;
    }
    tree.resize(n);
    for(int i=0;i<n-1;i++){
        int a,b;
        cin>>a>>b;
        a--;b--;
        tree[a].push_back(b);
        tree[b].push_back(a);
    }
    dfs(0);
    cout<<ans<<endl;
    return 0;
}