美团3.12笔试第5题

现在给你一颗树，每个树上的节点会被涂成黑色或白色。
现在定义好节点：
——对于白色的节点：若该节点没有子节点，或该节点子节点中至少有一个为黑色节点，则该节点是好节点
——对于黑色的节点：若该节点没有子节点，或该节点的所有子节点均为白色节点，则该节点是好节点

你的任务是找出这棵树上黑色的好节点和白色的好节点各有几个。

输入描述
第一行一个正整数n，表示这棵树共有n个节点，编号1到n。
第二行n个空格隔开的正整数，代表每个节点的颜色。0是白色，1是黑色。
接下来n个空格隔开的正整数，第 i 个正整数v表示节点 i 的父节点是v。其中，数字0表示这个节点是根。
1 <= n <= 10000

输出描述
一行两个正整数，以空格分开，第一个正整数表示白色的好节点有几个，第二个正整数表示黑色的好节点有几个。
样例输入
6
1 0 1 1 0 0
0 1 2 1 4 4
样例输出
3 2

#include<bits/stdc++.h>
using namespace std;
const int N = 10000+10;
int e[N],ne[N],h[N],color[N],idx;
int ans[2];//ans[0]是白色，ans[1]是黑色
void insert(int a,int b){//a->b
    e[idx] = b;
    ne[idx] = h[a];
    h[a] = idx++;
}
void dfs(int node){
    if(h[node] == -1){
        ans[color[node]]++;
        return;
    }
    int black = 0;
    for(int i = h[node]; i!=-1; i=ne[i]){
        int x = e[i];
        if(color[x]) black++;
        dfs(x);
    }
    if(color[node]==0 && black > 0) ans[0]++;
    else if(color[node]==1 && black == 0) ans[1]++;
}
void bfs(int root){
    queue<int> q;
    q.push(root);
    while(q.size()){
        int t = q.front();
        q.pop();
        if(h[t] == -1){
            ans[color[t]]++;
        }
        else{
            int black = 0;
            for(int i = h[t]; i!=-1; i=ne[i]){
                int x = e[i];
                if(color[x]) black++;
                    q.push(x);
            }
            if(color[t]==0 && black > 0) ans[0]++;
            else if(color[t]==1 && black == 0) ans[1]++;
        }
    }
}
int main(){
    memset(h,-1,sizeof h);
    int n;
    cin>>n;
    for(int i = 0;i < n;i++){
        cin>>color[i+1];
    }
    int root;
    for(int i = 1;i <= n;i++){
        int fa;
        cin>>fa;
        if(fa)
            insert(fa,i);
        else
            root = i;
    }
    bfs(root);
    cout<<ans[0]<<' '<<ans[1]<<endl;
    return 0;
}  

深搜和宽搜都可以哈