美团3.12笔试第五题

#include <vector>

using namespace std;

/*
6
1 0 1 1 0 0
0 1 2 1 4 4
对于白色的节点：若该节点没有子节点，或该节点子节点中至少有一个为黑色节点，则该节点是好节点
对于黑色的节点：若该节点没有子节点，或该节点的所有子节点均为白色节点，则该节点是好节点
*/

class Solution {
public:
 vector<int> getres(vector<int>& colors, vector<int>& parents) {
  int n = colors.size();
  vector<vector<int>> sons(n);
  for (int i = 1; i < n; i++) {
   sons[parents[i] - 1].push_back(i);
  }
  int blacks = 0, whites = 0;
  for (int i = 0; i < n; i++) {
   if (colors[i] == 0) {
    if (sons[i].empty()) {
     whites++;
    }
    else {
     for (int& num : sons[i]) {
      if (colors[num] == 1) {
       whites++;
       break;
      }
     }
    }
   }
   else {
    if (sons[i].empty()) {
     blacks++;
    }
    else {
     bool flag = false;
     for (int& num : sons[i]) {
      if (colors[num] == 1) {
       flag = true;
       break;
      }
     }
     if (flag == false) {
      blacks++;
     }
    }
   }
  }
  return { whites, blacks };
 }
};

int main() {
 int n;
 while (cin >> n) {
  vector<int> colors(n);
  vector<int> parents(n);
  for (int i = 0; i < n; i++) {
   cin >> colors[i];
  }
  for (int i = 0; i < n; i++) {
   cin >> parents[i];
  }
  Solution Sol;
  vector<int> res = Sol.getres(colors, parents);
  cout << res[0] << ' ' << res[1] << endl;
 }
 return 0;
}