#include "split.h"
#include <bits/stdc++.h>
using namespace std;
vector<vector<pair<int, int>>> adj;
vector<int> sz, label;
pair<int, int> g[3];
int N;
int DFS(int index, int last)
{
sz[index] = 1;
for (auto&node:adj[index])
{
if (node.second == last) continue;
node.first = DFS(node.second, index);
sz[index] += node.first;
}
for (auto&node : adj[index])
{
if (node.second != last) continue;
node.first = N - sz[index];
break;
}
return sz[index];
}
int lbl_a = 0;
void labelA(int index, int last)
{
if (lbl_a < g[0].first)
label[index] = g[0].second;
else
label[index] = g[2].second;
lbl_a++;
for (auto node : adj[index])
{
if (node.second == last) continue;
labelA(node.second, index);
}
}
int lbl_b = 0;
void labelB(int index, int last)
{
if (lbl_b < g[1].first)
label[index] = g[1].second;
else
label[index] = g[2].second;
lbl_b++;
for (auto node : adj[index])
{
if (node.second == last) continue;
labelB(node.second, index);
}
}
void circleDFS(int index)
{
if (lbl_a < g[0].first)
{
label[index] = g[0].second;
lbl_a++;
}
else if (lbl_b< g[1].first)
{
label[index] = g[1].second;
lbl_b++;
}
else
label[index] = g[2].second;
for (auto node : adj[index])
{
if (label[node.second]) continue;
circleDFS(node.second);
}
}
vector<int> find_split(int n, int a, int b, int c, vector<int> p, vector<int> q)
{
N = n;
const int M = p.size();
g[0] = {a, 1}; g[1] = {b, 2}; g[2] = {c, 3};
sort(g, g + 3);
adj = vector<vector<pair<int, int>>>(N);
sz = label = vector<int>(N);
for (int i = 0; i < M; i++)
{
adj[p[i]].push_back({0, q[i]});
adj[q[i]].push_back({0, p[i]});
}
if (M != N - 1)
{
circleDFS(0);
return label;
}
// Let a be the smallest
// if sz[i] >= a and (N - sz[i]) >= min(b, c) or opposite it can be done
// Find tree sizes
DFS(0, 0);
for (int i = 0; i < N; i++)
{
sort(adj[i].begin(), adj[i].end());
// Find the smallest subtree larger than g[0].first
if (adj[i].back().first < g[0].first) continue;
int l = 0, r = adj[i].size() - 1;
while (l < r)
{
int c = (l + r) / 2;
if (adj[i][c].first < g[0].first)
l = c + 1;
else
r = c;
}
int subtree = adj[i][l].first;
if (sz[i] - subtree < g[1].first) continue;
labelA(adj[i][l].second, i);
labelB(i, adj[i][l].second);
break;
}
return label;
}
