#include "beechtree.h"
#include <bits/stdc++.h>
using namespace std;
std::vector<int> beechtree(int N, int M, std::vector<int> P, std::vector<int> C) {
vector<int> sh(C);
sort(sh.begin(), sh.end());
sh.resize(unique(sh.begin(), sh.end()) - sh.begin());
M = size(sh);
for (int i = 0; i < N; ++i) {
C[i] = lower_bound(sh.begin(), sh.end(), C[i]) - sh.begin();
}
vector<int> ans(N, 1), isOk(N, 1), siz(N, 1);
vector<vector<int>> adj(N), sub(N);
for (int i = N - 1; i > 0; --i) {
adj[P[i]].push_back(i);
}
for (int v = N - 1; v >= 0; --v) {
static vector<int> cnt(M);
for (int to : adj[v]) {
cnt[C[to]] = 0;
isOk[v] &= isOk[to];
}
for (int to : adj[v]) {
siz[v] += siz[to];
isOk[v] &= ++cnt[C[to]] == 1;
}
ans[v] = isOk[v];
if (!isOk[v]) {
continue;
}
for (int to : adj[v]) {
sub[v].insert(sub[v].end(), sub[to].begin(), sub[to].end());
sub[to].clear();
}
static vector<int> last(N, -1);
last[v] = -1;
for (int x : sub[v]) {
last[x] = -1;
++cnt[C[x]];
}
// vector<int> ord = sub[v];
// sort(ord.begin(), ord.end(), [&](int x, int y) { return pair{cnt[C[x]], C[x]} > pair{cnt[C[y]], C[y]}; });
// for (int i = 0, j = 0, k = -1; i < size(ord); i = j) {
// int sum = 0;
// while (j < size(ord) && C[ord[j]] == C[ord[i]]) {
// int x = ord[j++];
// while (x != v) {
// x = P[x];
// if (last[x] == i) {
// break;
// } else if (last[x] != k) {
// ans[v] = 0;
// break;
// }
// sum += 1;
// last[x] = i;
// }
// if (!ans[v]) {
// break;
// }
// }
// if (!ans[v] || sum != j - i) {
// ans[v] = 0;
// break;
// }
// k = i;
// }
// if (!ans[v]) {
// sub[v].push_back(v);
// continue;
// }
vector<int> q;
set<array<int, 3>> st;
vector<int> tin(N);
st.insert({siz[v], tin[v] = 0, v});
int T = 0;
for (int i = 0; !st.empty() && ans[v]; ++i) {
q.push_back(st.rbegin()->begin()[2]);
st.erase(--st.end());
int x = q[i];
if (i > 0) {
int c = 0;
for (int j = 1; j < i; ++j) {
c += C[q[j]] == C[x];
}
if (P[x] != q[c]) {
ans[v] = 0;
break;
}
}
sort(adj[x].begin(), adj[x].end(), [&](int a, int b) { return siz[a] > siz[b]; });
for (int to : adj[x]) {
st.insert({siz[to], tin[to] = --T, to});
}
}
sub[v].push_back(v);
}
return ans;
}
