题目链接:http://codeforces.com/contest/1118/problem/F1
题意:现在有一棵树,树上的每个节点可以涂成红色、蓝色,或者不涂色,现在你可以删去一条边,将一棵树变成两棵树,两个树都只包含一种颜色。问这样的边有几个。
解题心得:dfs枚举得到以当前点为新的树的根节点分别包含了红色、蓝色节点的数量,该节点所在的子树包含了所有的红色并且没有蓝色,或者包含了所有的蓝色没有红色那么答案加1,一个dfs解决问题。
#include <bits/stdc++.h>
using namespace std;
const int maxn = 3e5+100;
int num[maxn];
vector <int> ve[maxn];//建立边
pair<int,int> P[maxn];//记录以该节点为根节点的树含有的两种颜色的个数
int n, cnt_1, cnt_2;
void init() {
scanf("%d", &n);
for(int i=1;i<=n;i++) {
scanf("%d", &num[i]);
if(num[i] == 1) cnt_1 ++;
else if(num[i] == 2) cnt_2++;
}
for(int i=1;i<n;i++) {
int a, b; scanf("%d%d",&a, &b);
ve[a].push_back(b);
ve[b].push_back(a);
}
}
int ans = 0;
pair <int, int> dfs(int pre, int pos) {
for(int i=0;i<ve[pos].size();i++) {
int Next = ve[pos][i];
if(Next == pre) continue;
pair<int,int> temp = dfs(pos, Next);
P[pos].first += temp.first;
P[pos].second += temp.second;
}
if(num[pos] == 1) {
P[pos].first ++;
} else if(num[pos] == 2){
P[pos].second ++;
}
if(P[pos].first == cnt_1 && P[pos].second == 0) ans++;
else if(P[pos].second == cnt_2 && P[pos].first == 0) ans++;
return P[pos];
}
int main() {
// freopen("1.in", "r", stdin);
init();
dfs(0, 1);
printf("%d\n", ans);
return 0;
}