#include <bits/stdc++.h> #include "beechtree.h" using namespace std; vector<int> b; vector<int> leafDist; vector<set<int>> vs; vector<vector<int>> adj; void DFS(int index, vector<int>&C) { b[index] = 1; if (adj[index].size() == 0) return; for (auto node : adj[index]) { DFS(node, C); leafDist[index] = max(leafDist[index], leafDist[node] + 1); } if (leafDist[index] == 1) { for (auto node : adj[index]) { if (vs[index].count(C[node])) b[index] = 0; vs[index].insert(C[node]); } } else b[index] = 0; } vector<int> beechtree(int N, int M, vector<int> P, vector<int> C) { leafDist = b = vector<int>(N); vs = vector<set<int>>(N); adj = vector<vector<int>>(N); for (int i = 1; i < N; i++) adj[P[i]].push_back(i); DFS(0, C); for (int i = 0; i < N; i++) { if (leafDist[i] != 2) continue; int index = -1; for (auto node : adj[i]) { if (leafDist[node] == 1) { if (index != -1) { index = -1; break; } else { index = node; } } } if (index == -1) continue; b[i] = 1; set<int> s; for (auto node : adj[i]) { if (s.count(C[node])) { b[i] = 0; break; } s.insert(C[node]); } for (auto n : vs[index]) { if (!s.count(n)) { b[i] = 0; break; } } } return b; }