#include <bits/stdc++.h>
#include "split.h"
template <class T> using Vec = std::vector<T>;
template <class F> struct RecLambda: private F {
explicit RecLambda(F&& f): F(std::forward<F>(f)) {}
template <class... Args> decltype(auto) operator () (Args&&... args) const {
return F::operator()(*this, std::forward<Args>(args)...);
}
};
struct DSU {
Vec<int> par;
DSU(const int n): par(n, -1) {}
int find(const int x) {
return par[x] < 0 ? x : par[x] = find(par[x]);
}
bool merge(int x, int y) {
x = find(x);
y = find(y);
if (x != y) {
par[x] = y;
return true;
}
return false;
}
};
Vec<int> find_split(int n, int a, int b, int c, Vec<int> x, Vec<int> y) {
std::array<int, 4> specified_size = {-1, a, b, c};
std::array<int, 3> size_order = {1, 2, 3};
std::sort(size_order.begin(), size_order.end(), [&](const int i, const int j) {
return specified_size[i] < specified_size[j];
});
const auto [small, medium, big] = size_order;
const int m = (int) x.size();
DSU dsu(n);
Vec<std::pair<int, int>> tree;
Vec<std::pair<int, int>> other;
for (int i = 0; i < m; ++i) {
const auto u = x[i];
const auto v = y[i];
if (dsu.merge(u, v)) {
tree.emplace_back(u, v);
} else {
other.emplace_back(u, v);
}
}
Vec<Vec<int>> graph(n);
for (const auto [u, v]: tree) {
graph[u].push_back(v);
graph[v].push_back(u);
}
Vec<int> subtree(n);
RecLambda([&](auto&& dfs, const int u, const int p) -> void {
subtree[u] = 1;
for (const auto v: graph[u]) {
if (v != p) {
dfs(v, u);
subtree[u] += subtree[v];
}
}
})(0, -1);
Vec<int> ret(n);
const auto write_tree = RecLambda([&](auto&& dfs, const int u, const int p, const int w, int& c) -> void {
if (c == 0) {
return;
}
ret[u] = w;
c -= 1;
for (const auto v: graph[u]) {
if (v != p) {
dfs(v, u, w, c);
}
}
});
bool done = false;
for (auto& [u, v]: tree) {
if (subtree[u] < subtree[v]) {
std::swap(u, v);
}
if (subtree[v] >= specified_size[small] and n - subtree[v] >= specified_size[medium]) {
write_tree(v, u, small, specified_size[small]);
write_tree(u, v, medium, specified_size[medium]);
done = true;
break;
}
if (subtree[v] >= specified_size[medium] and n - subtree[v] >= specified_size[small]) {
write_tree(u, v, small, specified_size[small]);
write_tree(v, u, medium, specified_size[medium]);
done = true;
break;
}
}
if (!done) {
const auto cent = RecLambda([&](auto&& dfs, const int u, const int p) -> int {
for (const auto v: graph[u]) {
if (v != p and subtree[v] * 2 > n) {
return dfs(v, u);
}
}
return u;
})(0, -1);
DSU split(n);
for (const auto [u, v]: tree) {
if (u != cent and v != cent) {
split.merge(u, v);
}
}
for (const auto [u, v]: other) {
if (u != cent and v != cent) {
split.merge(u, v);
const int r = split.find(u);
if (-split.par[r] >= specified_size[small]) {
for (const auto [u, v]: other) {
graph[u].push_back(v);
graph[v].push_back(u);
}
RecLambda([&](auto&& dfs, const int u) -> void {
if (specified_size[small] == 0 or ret[u] != 0 or split.find(u) != r) {
return;
}
ret[u] = small;
specified_size[small] -= 1;
for (const auto v: graph[u]) {
dfs(v);
}
})(r);
RecLambda([&](auto&& dfs, const int u) -> void {
if (specified_size[medium] == 0 or ret[u] != 0 or split.find(u) == r) {
return;
}
ret[u] = medium;
specified_size[medium] -= 1;
for (const auto v: graph[u]) {
dfs(v);
}
})(cent);
break;
}
}
}
}
if (std::any_of(ret.begin(), ret.end(), [&](const int x) { return x != 0; })) {
for (auto& x: ret) {
if (x == 0) {
x = big;
}
}
}
return ret;
}
