#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)...);
}
};
Vec<int> find_split(int n, int a, int b, int c, Vec<int> edge_p, Vec<int> edge_q) {
assert((int) edge_p.size() == n - 1);
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;
Vec<int> subtree(n);
Vec<Vec<int>> graph(n);
for (int i = 0; i < (int) edge_p.size(); ++i) {
graph[edge_p[i]].push_back(edge_q[i]);
graph[edge_q[i]].push_back(edge_p[i]);
}
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 = 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);
}
}
});
int X = specified_size[0];
int Y = specified_size[1];
for (int u = 0; u < n; ++u) {
const int len = (int) graph[u].size();
for (int i = 0; i < len; ++i) {
const int size = (subtree[graph[u][i]] > subtree[u] ? n - subtree[u] : subtree[graph[u][i]]);
if (size >= X and n - size >= Y) {
write(graph[u][i], u, small, X);
write(u, graph[u][i], medium, Y);
goto answer;
}
if (size >= Y and n - size >= X) {
write(u, graph[u][i], small, X);
write(graph[u][i], u, medium, Y);
goto answer;
}
}
}
answer:
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;
}
