#include "beechtree.h"
#include <bits/stdc++.h>
using namespace std;
using vi = vector<int>;
using vvi = vector<vi>;
vi par, colour, ans;
vvi adj;
bool permutation_exists(vi &perm, set<int> &avail) {
    if (avail.size() == 0) {
        bool valid = true;
        for (int i = 1; i < perm.size(); i++) {
            int fi = 0; for (int j = 1; j < i; j++) if (colour[perm[j]] == colour[perm[i]]) fi++;
            valid = valid && par[perm[i]] == perm[fi];
        }
        return valid;
    }
    set<int> new_avail = avail;
    for (auto &el : avail) {
        new_avail.erase(el);
        perm.push_back(el);
        if (permutation_exists(perm, new_avail)) return true;
        perm.pop_back();
        new_avail.insert(el);
    }
    return false;
}
void dfs(int u, int p, set<int> &st) {
    st.insert(u);
    for (int v : adj[u]) {
        if (v == p) continue;
        set<int> child_avail;
        dfs(v, u, child_avail);
        st.merge(child_avail);
    }
    vi perm = {u};
    st.erase(u);
    ans[u] = permutation_exists(perm, st) ? 1 : 0;
    st.insert(u);
}
vi beechtree(int N, int M, vi P, vi C) {
    par = P; colour = C;
    adj.assign(N, vi());
    for (int i = 1; i < N; i++) {
        adj[P[i]].push_back(i);
        adj[i].push_back(P[i]);
    }
    ans.assign(N, 0);
    set<int> avail;
    // dfs(0, 0, avail);
    
    int acount = 0, bcount = 0;
    set<int> acol, bcol;
    for (int i = 1; i < N; i++) {
        ans[i] = 1;
        if (P[i] == 0) {
            acount++;
            acol.insert(colour[i]);
        } else {
            bcount++;
            bcol.insert(colour[i]);
        }
    }
    if (acol.size() == acount && bcol.size() == 1 && acol.count(*bcol.begin())) {
        ans[0] = 1;
    } else {
        ans[0] = 0;
    }
    for (int i = 1; i < N; i++) {
        if (adj[i].size() == 1) {
            ans[i] = 1;
            continue;
        }
        set<int> childcol;
        for (auto &v : adj[i]) {
            if (v == P[i]) continue;
            childcol.insert(colour[v]);
        }
        ans[i] = childcol.size() != (adj[i].size() - 1) ? 1 : 0;
    }
    return ans;
}
| # | Verdict  | Execution time | Memory | Grader output | 
|---|
| Fetching results... | 
| # | Verdict  | Execution time | Memory | Grader output | 
|---|
| Fetching results... | 
| # | Verdict  | Execution time | Memory | Grader output | 
|---|
| Fetching results... | 
| # | Verdict  | Execution time | Memory | Grader output | 
|---|
| Fetching results... | 
| # | Verdict  | Execution time | Memory | Grader output | 
|---|
| Fetching results... | 
| # | Verdict  | Execution time | Memory | Grader output | 
|---|
| Fetching results... | 
| # | Verdict  | Execution time | Memory | Grader output | 
|---|
| Fetching results... | 
| # | Verdict  | Execution time | Memory | Grader output | 
|---|
| Fetching results... | 
| # | Verdict  | Execution time | Memory | Grader output | 
|---|
| Fetching results... | 
| # | Verdict  | Execution time | Memory | Grader output | 
|---|
| Fetching results... |