#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;
        static 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;
}