#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;
vector<int> dfs(int u){
int nn = 1 + pov[u].size();
if(nn == 1){return {};}
unordered_set<int> contains;
vector<vector<int>> a(nn);
vector<pair<int,int>> sub(nn);
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);
sub[i+1] = {-a[i+1].size(), i+1};
if(!beautiful[v]){beautiful[u] = 0;}
}
if(!beautiful[u]){return {};}
sub[0] = {-nn+1, 0};
sort(a[0].begin(), a[0].end());
sort(sub.begin(), sub.end());
for(int i = 1; i<nn;i++){
int l = 0;
int c_i = sub[i].second, prev_i = sub[i-1].second;
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 {};
}
}
}
return a[0];
}
vector<int> beechtree(int N, int M, vector<int> P, vector<int> C)
{
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 = 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';
// vector<int> 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';
// }
