# | Time | Username | Problem | Language | Result | Execution time | Memory |
---|---|---|---|---|---|---|---|
1232749 | antonn | Beech Tree (IOI23_beechtree) | C++20 | 0 ms | 0 KiB |
#include "beechtree.h"
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
const int N = 2e5 + 7;
int n, m;
vector<int> g[N];
int col[N], ok[N], par[N];
set<int> s[N];
int sz[N], d[N];
void dfs(int u) {
for (auto v : g[u]) {
dfs(v);
}
if (g[u].size() == 0) {
ok[u] = 1;
continue;
}
if (u != 0) {
set<int> children;
for (auto v : g[u]) children.insert(col[v]);
if (children.size() == g[u].size()) ok[u] = 1;
} else {
set<int> children;
for (auto v : g[u]) children.insert(col[v]);
if (children.size() != g[u].size()) ok[u] = 0;
for (auto v : g[u]) if (ok[v] == 0) ok[u] = 0;
}
}
vector<int> beechtree(int N, int M, vector<int> P, vector<int> C) {
n = N;
m = M;
for (int i = 0; i < n; ++i) {
par[i] = P[i];
if (i > 0) {
g[P[i]].push_back(i);
}
}
map<int, int> f;
ok[0] = 1;
for (int i = 0; i < n; ++i) {
col[i] = C[i];
f[col[i]]++;
}
int cnt = 0;
for (auto x : f) {
if (x.second > 2) {
ok[0] = 0;
} else if (x.second == 2) {
++cnt;
}
}
if (cnt > 1) ok[0] = 0;
dfs(0);
vector<int> ret(n);
for (int i = 0; i < n; ++i) ret[i] = ok[i];
return ret;
}