#include "beechtree.h"
#include <algorithm>
#include <array>
#include <vector>
std::vector<int> sol;
std::vector<std::vector<std::pair<int, int>>> children;
std::vector<int> parent;
std::vector<int> color;
std::vector<int> DFS(int c) {
std::vector<int> nodes = {c};
for (auto [next, color] : children[c]) {
std::vector<int> next_nodes = DFS(next);
for (auto a : next_nodes) nodes.push_back(a);
}
std::sort(nodes.begin(), nodes.end());
do {
if (nodes[0] != c) continue;
std::array<int, 502> colors; colors.fill(0);
bool good = true;
for (int i = 1; i < nodes.size(); i++) {
if (nodes[colors[color[nodes[i]]]] != parent[nodes[i]]) {
good = false;
break;
}
colors[color[nodes[i]]]++;
}
if (good) {
sol[c] = 1;
break;
}
} while(std::next_permutation(nodes.begin(), nodes.end()));
return nodes;
}
std::vector<int> beechtree(int N, int M, std::vector<int> P, std::vector<int> C)
{
sol.assign(N, 0);
parent = P;
color = C;
children.assign(N, {});
for (int i = 1; i < N; i++) children[P[i]].push_back({i, C[i]});
DFS(0);
return sol;
}
