#include "beechtree.h"
#include <iostream>
#include <vector>
#include <algorithm>
#include <unordered_map>
#include <unordered_set>
using namespace std;
int n, m;
vector<vector<int>> pov;
vector<int> beautiful, c;
int cur_d;
vector<int> dfs(int u) {
int nn = 1 + pov[u].size();
if (nn == 1) { cur_d = 0; return {}; }
unordered_set<int> contains;
vector<vector<int>> a(nn);
vector<pair<int, pair<int, int>>> sub(nn);
int max_d = 0;
for (int i = 0; i < nn - 1; i++) {
int v = pov[u][i];
if (contains.count(c[v])) {
beautiful[u] = 0;
}
contains.insert(c[v]);
a[0].push_back(c[v]);
a[i + 1] = dfs(v);
max_d = max(cur_d, max_d);
sub[i + 1] = { -(int)(a[i + 1].size()), {i + 1, cur_d}};
if (!beautiful[v]) { beautiful[u] = 0; }
}
if (!beautiful[u]) { return {}; }
max_d++;
sub[0] = { -nn + 1, {0,max_d} };
sort(a[0].begin(), a[0].end());
sort(sub.begin(), sub.end());
for (int i = 1; i < nn; i++) {
int l = 0;
if (sub[i].second.second > sub[i - 1].second.second) { beautiful[u] = 0; return {}; }
int c_i = sub[i].second.first, prev_i = sub[i - 1].second.first;
for (int j = 0; j < a[c_i].size(); j++) {
while (a[prev_i][l] < a[c_i][j]) {
l++;
if (l >= a[prev_i].size()) {
l--; break;
}
}
if (a[prev_i][l] != a[c_i][j]) {
beautiful[u] = 0; return {};
}
}
}
cur_d = max_d;
return a[0];
}
vector<int> beechtree(int N, int M, vector<int> P, vector<int> C)
{
pov.clear();
beautiful.clear();
c.clear();
n = N;
c = C;
pov.resize(n);
unordered_map<int, int> comp;
for (int i = 1; i < n; i++) {
pov[P[i]].push_back(i);
if (!comp.count(c[i])) { comp[c[i]] = comp.size(); }
c[i] = comp[c[i]];
}
m = comp.size();
beautiful.resize(n, 1);
dfs(0);
return beautiful;
}
// int main() {
// vector<int> ans = beechtree(4, 2, { -1, 0, 0, 0 }, { 0, 1, 1, 2 });
// for (int i = 0; i < ans.size(); i++) {
// cout << ans[i] << ' ';
// }
// cout << '\n';
// ans.clear();
// ans = beechtree(18, 3, { -1, 0, 0, 0, 1, 1, 1, 2, 2, 3, 3, 3, 4, 4, 5, 10, 11, 11 }, { 0, 1, 2, 3, 1, 2, 3, 1, 3, 3, 2, 1, 1, 2, 2, 1, 2, 3 });
// for (int i = 0; i < ans.size(); i++) {
// cout << ans[i] << ' ';
// }
// cout << '\n';
// ans.clear();
// ans = beechtree(7, 2, { -1, 0, 1, 1, 0, 4, 5 }, { 0, 1, 1, 2, 2, 1, 1 });
// for (int i = 0; i < ans.size(); i++) {
// cout << ans[i] << ' ';
// }
// cout << '\n';
// }
