#include "beechtree.h"
#include <bits/stdc++.h>
using namespace std;
/*
// sub 3 doesn't work because it doesn't consider colour of direct children of node 0
// unordered_map<int, int> dfs(vector<vector<int>> &children, vector<int> &P, vector<int> &C, vector<int> &res, int cur) {
// unordered_map<int, int> curM;
// for (auto &i : children[cur]) {
// auto um = dfs(children, P, C, res, i);
// um[C[i]] = 1;
// for (auto &[colour, cnt] : um) {
// if (curM[colour]) return curM;
// }
// }
// res[cur] = 1;
// return curM;
// }
// std::vector<int> beechtree(int N, int M, std::vector<int> P, std::vector<int> C)
// {
// vector<vector<int>> children(N);
// for (int i = 1; i < N; ++i) children[P[i]].push_back(i);
// vector<int> res(N, 0);
// vector<unordered_map<int, int>> maps;
// bool ok = true;
// for (auto &child : children[0]) {
// if (ok) maps.push_back(dfs(children, P, C, res, child));
// else dfs(children, P, C, res, child);
// if (res[child] == 0) ok = false;
// }
// if (ok) {
// sort(maps.begin(), maps.end(), [](const unordered_map<int, int> &a, const unordered_map<int, int> &b) {
// return a.size() < b.size();
// });
// int sz = maps.size();
// for (int i = 1; i < sz; ++i) {
// for (auto &[colour, cnt] : maps[i]) {
// if (!maps[i - 1][colour]) {
// ok = false;
// break;
// }
// }
// if (!ok) break;
// }
// }
// if (ok) res[0] = 1;
// return res;
// }
*/
unordered_map<int, int> FAIL;
unordered_map<int, int> dfs(vector<vector<int>> &children, vector<int> &P, vector<int> &C, vector<int> &res, int cur) {
unordered_map<int, int> curM, childM;
int depth = 0;
bool fail = false;
for (auto &i : children[cur]) {
auto um = dfs(children, P, C, res, i);
if (um[-1]) fail = true;
um.erase(-1);
if (curM[C[i]]) fail = true;
if (!fail) {
++curM[C[i]];
if (!childM.empty() && !um.empty()) fail = true;
else if (!um.empty()) childM = um;
}
}
if (fail) return FAIL;
if (!childM.empty()) depth = 2;
else if (!curM.empty()) depth = 1;
if (depth <= 1) {
res[cur] = 1;
return curM;
}
for (auto &[colour, cnt] : childM) {
if (!curM[colour]) return FAIL;
}
res[cur] = 1;
return FAIL;
}
std::vector<int> beechtree(int N, int M, std::vector<int> P, std::vector<int> C)
{
vector<vector<int>> children(N);
for (int i = 1; i < N; ++i) children[P[i]].push_back(i);
vector<int> res(N, 0);
FAIL[-1] = 1;
dfs(children, P, C, res, 0);
return res;
}
