#include "split.h"
#include <bits/stdc++.h>
using namespace std;
template<typename T>
using vec = vector<T>;
using ll = long long;
using vll = vec<ll>;
using vvll = vec<vll>;
using pll = pair<ll,ll>;
#define pb push_back
#define dbg(x) cerr << #x << ": " << x << endl
#define cerr if(0) cout
/* type shit */
vll subtree_size;
vec<int> color;
vec<bool> vis, moveable, moveable2, moved, is_in_centroid_subtree, is_in_centroid2_subtree;
vec<pll> days;
vll dayOrder(3);
vll parent;
vvll adj, dt_adj;
ll N;
ll centroid = -1;
ll centroid2 = -1;
void dfs(ll u, ll p) {
parent[u] = p;
for(ll &v : adj[u]) {
if(v == p or vis[v]) continue;
vis[v] = 1;
dt_adj[u].pb(v);
dt_adj[v].pb(u);
dfs(v,u);
}
}
void dfs_tree(ll u, ll p) {
subtree_size[u] = 1;
for(ll &v : dt_adj[u]) {
if(v == p) continue;
dfs_tree(v,u);
subtree_size[u] += subtree_size[v];
}
if(centroid == -1 and subtree_size[u] >= days[0].first) {
centroid = u;
}
if(centroid2 == -1 and subtree_size[u] >= days[1].first) {
centroid2 = u;
}
}
void dfs_cassign(ll u, ll p) {
// just to check if things in centroid subtree
is_in_centroid_subtree[u] = is_in_centroid_subtree[p];
is_in_centroid2_subtree[u] = is_in_centroid2_subtree[p];
if(u == centroid) is_in_centroid_subtree[u] = 1;
if(u == centroid2) is_in_centroid2_subtree[u] = 1;
for(ll &v : dt_adj[u]) {
if(v == p) continue;
dfs_cassign(v,u);
}
}
bool check_moveable(ll x, ll ct) {
// returns 1 if we can bfs from here to some node that is not in centroid subtree
queue<ll> q;
q.push(x);
vis[x] = 1;
bool ans = 0;
while(!q.empty()) {
ll u = q.front();
q.pop();
for(ll &v : adj[u]) {
if(vis[v] or v == ct) continue;
vis[v] = 1;
if(!is_in_centroid_subtree[v]) ans = 1;
q.push(v);
}
}
return ans;
}
void dfs_color_first(ll u, ll p, ll c) {
if(days[dayOrder[c-1]].first == 0) return;
color[u] = c;
days[dayOrder[c-1]].first--;
for(ll &v : dt_adj[u]) {
if(v == p or moved[v]) continue;
dfs_color_first(v,u,c);
}
}
void bfs_color_second(ll x, ll c) {
queue<ll> q;
q.push(x);
color[x] = c;
days[dayOrder[c-1]].first--;
if(days[dayOrder[c-1]].first == 0) return;
while(!q.empty()) {
ll u = q.front();
q.pop();
for(ll &v : adj[u]) {
if(color[v] != -1) continue;
color[v] = c;
days[dayOrder[c-1]].first--;
if(days[dayOrder[c-1]].first == 0) return;
q.push(v);
}
}
}
vector<int> find_split(int n, int a, int b, int c, vector<int> p, vector<int> q) {
N = n;
ll m = p.size();
adj.resize(n);
dt_adj.resize(n);
parent.assign(n,-1);
vis.assign(n,0);
moveable.assign(n,0);
moveable2.assign(n,0);
moved.assign(n,0);
is_in_centroid_subtree.assign(n,0);
is_in_centroid2_subtree.assign(n,0);
for(int i = 0; i < m; i++) {
adj[p[i]].pb(q[i]);
adj[q[i]].pb(p[i]);
}
cerr << "Adj built" << endl;
days = {{a,1}, {b,2}, {c,3}};
sort(days.begin(),days.end());
for(int i = 0; i < 3; i++) dayOrder[days[i].second - 1] = i;
cerr << "Days: ";
for(auto [a,b] : days) cerr << "(" << a << ", " << b << "), ";
cerr << endl;
subtree_size.assign(n,0);
color.assign(n,-1);
vis[0] = 1;
dfs(0,-1);
cerr << "DFS tree built" << endl;
// cerr << "DT adj: " << endl;
// for(int i = 0; i < n; i++) {
// cerr << i << ": ";
// for(ll &x : dt_adj[i]) cerr << x << " ";
// cerr << endl;
// }
dfs_tree(0,-1);
if(centroid == -1 and centroid2 == -1) {
vec<int> ans(n,0);
return ans;
}
// cerr << "Subtree sizes: ";
// for(ll &x : subtree_size) cerr << x << " ";
cerr << endl;
cerr << "Centroid found: " << centroid << endl;
cerr << "Centroid2: " << centroid2 << endl;
dfs_cassign(0,0);
cerr << "Centroidness Assigned" << endl;
vis.assign(n,0);
if(centroid != -1) {
for(ll &x : dt_adj[centroid]) {
// for each child of centroid check if we can do surgery
if(x == parent[centroid] or vis[x]) continue;
vis[x] = 1;
moveable[x] = check_moveable(x, centroid);
if(moveable[x]) cerr << x << " is moveable (1)" << endl;
}
}
cerr << "Movability (1) complete" << endl;
vis.assign(n,0);
if(centroid2 != -1) {
for(ll &x : dt_adj[centroid2]) {
// for each child of centroid check if we can do surgery
if(x == parent[centroid2] or vis[x]) continue;
vis[x] = 1;
moveable2[x] = check_moveable(x, centroid2);
if(moveable2[x]) cerr << x << " is moveable (2) " << endl;
}
}
cerr << "Movability (2) complete" << endl;
// CASE: A fits in centroid subtree, B fits above centroid
ll c_st_size = -1;
ll rest = -1;
if(centroid != -1) {
c_st_size = subtree_size[centroid];
rest = n - c_st_size;
for(ll &x : dt_adj[centroid]) {
if(rest >= days[1].first and c_st_size >= days[0].first) break;
if(x == parent[centroid] or not moveable[x]) continue;
// assuming no heuristic and you just move whatever
if(c_st_size - subtree_size[x] >= days[0].first) { // move x on top
moved[x] = 1; // actually this represents te whole subtree
cerr << "Moved " << x << endl;
c_st_size -= subtree_size[x];
rest += subtree_size[x];
}
}
// did this case work?
if(c_st_size >= days[0].first and rest >= days[1].first) {
// yay! finish
dfs_color_first(centroid, parent[centroid], days[0].second);
bfs_color_second(parent[centroid], days[1].second);
for(int &x : color) x = (x != -1 ? x : days[2].second);
return color;
}
}
cerr << "Case A failed, moving on..." << endl;
// CASE: B fits in centroid subtree, A fits above centroid
if(centroid2 != -1) {
c_st_size = subtree_size[centroid2];
rest = n - c_st_size;
moved.assign(n,0);
for(ll &x : dt_adj[centroid2]) {
if(rest >= days[0].first and c_st_size >= days[1].first) break;
if(x == parent[centroid2] or not moveable2[x]) continue;
if(c_st_size - subtree_size[x] >= days[1].first) {
moved[x] = 1;
cerr << "Moved " << x << endl;
c_st_size -= subtree_size[x];
rest += subtree_size[x];
}
}
// dbg(c_st_size);
// dbg(rest);
if(c_st_size >= days[1].first and rest >= days[0].first) {
dfs_color_first(centroid2,parent[centroid2],days[1].second);
bfs_color_second(parent[centroid2],days[0].second);
cerr << "Sanity check: " << days[0].first << " " << days[1].first << endl;
for(int &x : color) x = (x != -1 ? x : days[2].second);
return color;
}
}
color.assign(n,0);
return color;
}
