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

题目链接：http://codeforces.com/contest/1118/problem/F1

题意：现在有一棵树，树上的每个节点可以涂成红色、蓝色，或者不涂色，现在你可以删去一条边，将一棵树变成两棵树，两个树都只包含一种颜色。问这样的边有几个。

解题心得：dfs枚举得到以当前点为新的树的根节点分别包含了红色、蓝色节点的数量，该节点所在的子树包含了所有的红色并且没有蓝色，或者包含了所有的蓝色没有红色那么答案加１，一个dfs解决问题。

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