#include<bits/stdc++.h>
#include <utility>
#include "beechtree.h"
#pragma GCC optimize("Ofast,unroll-loops")
#pragma GCC target("avx2,bmi,bmi2,popcnt,lzcnt")
using namespace std;
using g_t = vector<vector<int>>;
struct dVal {
unordered_map<int, int> de;
list<set<pair<int, int>>> se;
int seSi;
};
struct unordered_mapsVal {
map<int, int> maDe;
unordered_map<int, unordered_map<int, map<int, int>>> maCDe;
unordered_map<int, map<int, int>> maCat;
map<int, int> miDe;
unordered_map<int, unordered_map<int, map<int, int>>> miCDe;
map<int, map<int, int>> miCat;
};
struct sdVal {
vector<pair<int, int>> edg;
unordered_map<int, unordered_mapsVal> maV;
};
struct BFSSorting {
int n, actDe, actC, pDe, pC, cat, pO;
};
bool operator<(const BFSSorting &l, const BFSSorting &r) {
return tie(l.actDe, l.pDe, l.actC, l.pC, l.cat, r.pO) < tie(r.actDe, r.pDe, r.actC, r.pC, r.cat, l.pO);
}
vector<int> de;
vector<BFSSorting> de2;
g_t g;
vector<bool> po;
vector<int> c;
dVal dfs(int n) {
unordered_map<int, int> r;
unordered_map<int, vector<pair<int, int>>> ve;
list<set<pair<int, int>>> se;
unordered_map<int, unordered_map<int, int>> OmiR;
int seSi = 0;
for (auto e: g[n]) {
if (r.find(c[e]) != r.end()) po[n] = false;
r[c[e]] = 1;
ve[c[e]] = {};
}
for (auto e: g[n]) {
auto [re, li, ssi] = dfs(e);
if (!po[e]) {
po[n] = false;
}
if (!po[n]) continue;
for (auto &[co, d]: re) {
if (r.find(co) == r.end()) {
po[n] = false;
break;
}
if (co == c[e]) r[c[e]] = d + 1;
ve[co].emplace_back(d, e);
}
if (!po[n]) continue;
if (ssi > seSi) {
swap(se, li);
swap(ssi, seSi);
}
auto itS = se.begin();
for (auto &it: li) {
while (itS != se.end() && itS->size() <= it.size()) {
for (auto cos: *itS) {
if (it.erase(cos) == 0) {
po[n] = false;
break;
}
}
if (!po[n]) break;
itS++;
}
if (!po[n]) break;
if (!it.empty()) {
auto act = se.insert(itS, set<pair<int, int>>());
for (auto cos: it) {
if (itS->erase(cos) == 0) {
po[n] = false;
break;
}
act->insert(cos);
}
}
if (itS != se.end() && itS->empty()) itS++;
if (!po[n]) break;
}
}
if (!po[n]) return {};
se.emplace_back();
for (auto &[co, vec]: ve) {
bool dD = true;
for (auto [d, e]: vec) {
if (d > r[co]) {
po[n] = false;
return {};
} else if (d == r[co]) {
dD = false;
break;
}
}
if (dD) {
se.back().insert({co, r[co]});
seSi++;
}
}
if (se.back().empty()) se.pop_back();
de[n] = seSi;
return {move(r), move(se), seSi};
}
sdVal sdfs(int n, vector<bool> &vis) {
if (vis[n]) return {};
vis[n] = true;
sdVal res;
for (auto e: g[n]) {
auto re = sdfs(e, vis);
if (!po[e]) po[n] = false;
if (!po[n]) continue;
if (re.edg.size() > res.edg.size()) swap(re, res);
re.edg.emplace_back(e, n);
for (auto &ed: re.edg) {
res.edg.push_back(ed);
auto &ch = de2[ed.first];
auto &miDe = res.maV[ch.actC].miDe;
auto up = miDe.upper_bound(ch.pDe);
if (up != miDe.end() && up->second < ch.actDe) {
po[n] = false;
break;
}
if (miDe.find(ch.pDe) == miDe.end()) miDe[ch.pDe] = ch.actDe;
miDe[ch.pDe] = min(miDe[ch.pDe], ch.actDe);
auto &miCDe = res.maV[ch.actC].miCDe[ch.pC][ch.pDe];
up = miCDe.upper_bound(de2[ed.second].pDe);
if (up != miCDe.end() && up->second < ch.actDe) {
po[n] = false;
break;
}
if (miCDe.find(de2[ed.second].pDe) == miCDe.end()) miCDe[de2[ed.second].pDe] = ch.actDe;
miCDe[de2[ed.second].pDe] = min(miCDe[de2[ed.second].pDe], ch.actDe);
auto &miCat = res.maV[ch.actC].miCat[de2[ed.second].cat];
up = miCat.upper_bound(-ch.pO);
if (up != miCat.end() && up->second < ch.actDe) {
po[n] = false;
break;
}
if (miCat.find(-ch.pO) == miCat.end())cmiCat[-ch.pO] = ch.actDe;
miCat[-ch.pO] = min(miCat[-ch.pO], ch.actDe);
auto &maDe = res.maV[ch.actC].maDe;
up = maDe.upper_bound(-ch.pDe);
if (up !=maDe.end() && up->second > ch.actDe) {
po[n] = false;
break;
}
if (maDe.find(-ch.pDe) == maDe.end())cmaDe[-ch.pDe] = ch.actDe;
maDe[-ch.pDe] = max(maDe[-ch.pDe], ch.actDe);
auto &maCDe = res.maV[ch.actC].maCDe[ch.pC][ch.pDe];
up = maCDe.upper_bound(-de2[ed.second].pDe);
if (up != maCDe.end() && up->second > ch.actDe) {
po[n] = false;
break;
}
if (maCDe.find(-de2[ed.second].pDe) == maCDe.end())cmaCDe[-de2[ed.second].pDe] = ch.actDe;
maCDe[-de2[ed.second].pDe] = max(maCDe[-de2[ed.second].pDe], ch.actDe);
auto &maCat = res.maV[ch.actC].maCat[de2[ed.second].cat];
up = maCat.upper_bound(ch.pO);
if (up != maCat.end() && up->second > ch.actDe) {
po[n] = false;
break;
}
if (maCat.find(ch.pO) == maCat.end())cmaCat[ch.pO] = ch.actDe;
maCat[ch.pO] = max(maCat[ch.pO], ch.actDe);
}
}
return std::move(res);
}
void bfs(int n, vector<bool> &vi) {
int d = 0, col = 0;
BFSSorting la{n, de[n], c[n], de[n], c[n], 0, 1};
priority_queue<BFSSorting> pq;
pq.push({n, de[n], c[n], de[n], c[n], 0, 1});
while (!pq.empty()) {
auto bO = pq.top();
pq.pop();
if (vi[bO.n]) continue;
vi[bO.n] = true;
if (bO.actDe != la.actDe || bO.pDe != la.pDe || bO.pO != la.pO) d++;
if (bO.actC != la.actC || bO.pC != la.pC || bO.cat != la.cat || bO.actDe != la.actDe || bO.pDe != la.pDe) col++;
de2[bO.n] = la = {bO.n, bO.actDe, bO.actC, bO.pDe, bO.pC, col, d};
for (auto e: g[bO.n]) pq.push({e, de[e], c[e], de[bO.n], c[bO.n], de2[bO.n].cat, de2[bO.n].pO});
}
}
vector<int> beechtree(int N, int M, vector<int> P, vector<int> C) {
g = g_t(N);
for (int i = 1; i < N; ++i) g[P[i]].push_back(i);
de = vector<int>(N);
de2 = vector<BFSSorting>(N);
po= vector<bool> (N, true);
c = move(C);
dfs(0);
vector<bool> vis(N, false);
for (int i = 0; i < N; ++i) if (po[i]) bfs(i, vis);
fill(vis.begin(), vis.end(), false);
for (int i = 0; i < N; ++i) if (po[i]) sdfs(i, vis);}
return {po.begin(), po.end()};
}
Compilation message
beechtree.cpp: In function 'sdVal sdfs(int, std::vector<bool>&)':
beechtree.cpp:160:42: error: 'cmiCat' was not declared in this scope; did you mean 'miCat'?
160 | if (miCat.find(-ch.pO) == miCat.end())cmiCat[-ch.pO] = ch.actDe;
| ^~~~~~
| miCat
beechtree.cpp:168:41: error: 'cmaDe' was not declared in this scope; did you mean 'maDe'?
168 | if (maDe.find(-ch.pDe) == maDe.end())cmaDe[-ch.pDe] = ch.actDe;
| ^~~~~
| maDe
beechtree.cpp:176:55: error: 'cmaCDe' was not declared in this scope; did you mean 'maCDe'?
176 | if (maCDe.find(-de2[ed.second].pDe) == maCDe.end())cmaCDe[-de2[ed.second].pDe] = ch.actDe;
| ^~~~~~
| maCDe
beechtree.cpp:184:41: error: 'cmaCat' was not declared in this scope; did you mean 'maCat'?
184 | if (maCat.find(ch.pO) == maCat.end())cmaCat[ch.pO] = ch.actDe;
| ^~~~~~
| maCat
beechtree.cpp:188:18: warning: moving a local object in a return statement prevents copy elision [-Wpessimizing-move]
188 | return std::move(res);
| ~~~~~~~~~^~~~~
beechtree.cpp:188:18: note: remove 'std::move' call
beechtree.cpp: In function 'std::vector<int> beechtree(int, int, std::vector<int>, std::vector<int>)':
beechtree.cpp:219:54: warning: no return statement in function returning non-void [-Wreturn-type]
219 | for (int i = 0; i < N; ++i) if (po[i]) sdfs(i, vis);}
| ^
beechtree.cpp: At global scope:
beechtree.cpp:220:2: error: expected unqualified-id before 'return'
220 | return {po.begin(), po.end()};
| ^~~~~~
beechtree.cpp:221:1: error: expected declaration before '}' token
221 | }
| ^