Codeforces Round #540 (Div. 3)--1118F1 - Tree Cutting (Easy Version)

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=,blue=;
int ans=;

pair<int,int> dfs(int v,int p=-){
    int r=(color[v]==);
    int b=(color[v]==);
    for(int i=;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==);//只有之前经过的点中包含了全部的红或者蓝点才可以
            ans+=(tmp.first==&&tmp.second==blue);
            r+=tmp.first;//加上之前点中的红和蓝点
            b+=tmp.second;
        }
    }
    return make_pair(r,b);
}

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