#include "split.h"
#include <vector>
#include <cassert>
#include <cmath>
#include <algorithm>
#include <queue>
#include <iostream>
using ll = long long;
#define MIN(a, b) (((a) < (b)) ? (a) : (b))
#define MAX(a, b) (((a) < (b)) ? (b) : (a))
std::vector<std::vector<int>> g;
std::vector<std::vector<int>> tree;
std::vector<int> level;
std::vector<int> sz;
void dfs(int node) {
sz[node] = 1;
for(int h = 0; h < g[node].size(); h++) {
int to = g[node][h];
if(level[to] == 0) {
level[to] = level[node] + 1;
dfs(to);
tree[to].push_back(node);
tree[node].push_back(to);
sz[node] += sz[to];
}
}
}
int find_centroid(int node, int target) {
for(int h = 0; h < tree[node].size(); h++) {
int to = tree[node][h];
if(sz[to] < sz[node] && target <= sz[to])
return find_centroid(to, target);
}
return node;
}
void extract(int node, int parent, std::vector<int> &ord) {
for(int h = 0; h < tree[node].size(); h++) {
int to = tree[node][h];
if(to != parent)
extract(to, node, ord);
}
ord.push_back(node);
}
void mark(int node, std::vector<int> &seen, std::vector<int> &sol, int &scount, int color) {
if(scount == 0)
return ;
sol[node] = color;
scount--;
for(int h = 0; h < g[node].size(); h++) {
int to = g[node][h];
if(sol[to] == 3 && seen[to] == 1)
mark(to, seen, sol, scount, color);
}
}
std::vector<int> find_split(int n, int a, int b, int c, std::vector<int> p, std::vector<int> q) {
g.resize(n);
tree.resize(n);
level.resize(n);
sz.resize(n);
int id1 = 1, id2 = 2, id3 = 3;
if(b < a) {
std::swap(id1, id2);
std::swap(a, b);
}
if(c < a) {
std::swap(id1, id3);
std::swap(a, c);
}
if(c < b) {
std::swap(id2, id3);
std::swap(b, c);
}
for(int i = 0; i < p.size(); i++) {
int x = p[i];
int y = q[i];
g[x].push_back(y);
g[y].push_back(x);
}
level[0] = 1;
dfs(0);
std::vector<int> sol(n, 3);
int centroid = find_centroid(0, sz[0] / 2);
std::vector<int> mult(n);
std::vector<std::vector<int>> big;
for(int h = 0; h < tree[centroid].size(); h++) {
std::vector<int> aux;
extract(tree[centroid][h], centroid, aux);
big.push_back(aux);
for(int i = 0; i < aux.size(); i++)
mult[aux[i]] = h;
}
std::vector<int> used(n);
std::vector<std::vector<int>> bigg;
bigg.resize(big.size());
for(int i = 0; i < p.size(); i++) {
if(p[i] == centroid || q[i] == centroid)
continue;
if(mult[p[i]] != mult[q[i]]) {
bigg[mult[p[i]]].push_back(mult[q[i]]);
bigg[mult[q[i]]].push_back(mult[p[i]]);
}
}
for(int i = 0; i < big.size(); i++) {
if(used[i] == 0) {
std::vector<int> current;
std::queue<int> q;
q.push(i);
used[i] = 1;
while(0 < q.size() && current.size() < a) {
int node = q.front();
q.pop();
for(int h = 0; h < big[node].size(); h++)
current.push_back(big[node][h]);
for(int h = 0; h < bigg[node].size(); h++) {
int to = bigg[node][h];
if(used[to] == 0) {
used[to] = 1;
q.push(to);
}
}
}
if(current.size() < a)
continue;
else {
std::vector<int> seen(n, 0);
for(int i = 0; i < current.size(); i++)
seen[current[i]] = 1;
mark(current[0], seen, sol, a, 1);
seen = std::vector<int>(n, 1);
mark(centroid, seen, sol, b, 2);
for(int i = 0; i < sol.size(); i++)
if(sol[i] == 1)
sol[i] = id1;
else if(sol[i] == 2)
sol[i] = id2;
else if(sol[i] == 3)
sol[i] = id3;
return sol;
}
}
}
return std::vector<int>(n, 0);
}
