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

时间:2021-12-18 09:16:53

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 ;
}