#include "split.h" #include <bits/stdc++.h> using namespace std; vector<int> find_split(int n, int A, int B, int C, vector<int> P, vector<int> Q) { array<pair<int, int>, 3> att; att[0] = {A, 1}; att[1] = {B, 2}; att[2] = {C, 3}; sort(att.begin(), att.end()); A = att[0].first; B = att[1].first; C = att[2].first; vector<int> sz(n, 1); vector<vector<int>> adj(n); vector<vector<int>> anc(n); vector<int> vis(n, 0); vector<int> hgh(n); vector<int> dpth(n, 0); for (int i = 0; i < P.size(); i++) { adj[P[i]].push_back(Q[i]); adj[Q[i]].push_back(P[i]); } function<int(int, int)> lca = [&](int a, int b) { if (dpth[a] < dpth[b]) swap(a,b); int h = anc[a].size()-1; while (dpth[a] > dpth[b]) { while (h > 0 && dpth[anc[a][h]] < dpth[b]) h--; a = anc[a][h]; h = min((size_t)h, anc[a].size()-1); } if (a == b) return a; while (a != b) { while (h > 0 && anc[a][h] == anc[b][h]) h--; a = anc[a][h]; b = anc[b][h]; h = min((size_t)h, anc[a].size()-1); } return anc[a][0]; }; int base = 0; function<void(int, int)> dfs = [&](int a, int p) { if (base > 0) return; int h = 0; while (p != -1 && anc[p].size() > h) { anc[a].push_back(anc[p][h]); h++; } bool ok = 1; vis[a] = 1; hgh[a] = dpth[a]; for (int x : adj[a]) { if (x == p) continue; if (vis[x]) { hgh[a] = min(hgh[a], dpth[lca(a, x)]); continue; } anc[x].push_back(a); dpth[x] = dpth[a]+1; dfs(x, a); sz[a] += sz[x]; if (sz[x] >= A) { ok = 0; } } if (sz[a] < A) ok = 0; if (ok) { int aus = n-sz[a]; base = a; for (int x : adj[a]) { if (anc[x][0] != a) continue; if (aus >= B) break; if (hgh[x] < dpth[a]) aus += sz[x]; anc[x][0] = p; } if (aus < B) base = -1; } }; dfs(0, -1); if (base <= 0) return vector<int> (n, 0); vis = vector<int> (n, -1); int ja = 0, jb = 0; function<void(int, int, int)> dfsa = [&](int a, int p, int sign) { if (ja < A) { vis[a] = sign; ja++; } else return; for (int x : adj[a]) { if (x == p) continue; if (x == 0 || anc[x][0] != a) continue; dfsa(x, a, sign); } }; dfsa(base, anc[base][0], att[0].second); function<void(int, int)> dfsb = [&](int a, int sign) { vis[a] = 0; if (jb < B) { vis[a] = sign; jb++; } else return; for (int x : adj[a]) { if (vis[x] >= 0) continue; dfsb(x, sign); } }; dfsb(anc[base][0], att[1].second); for (int i = 0; i < n; i++) { if (vis[i] <= 0) vis[i] = att[2].second; } return vis; }