#include "split.h"
#include <algorithm>
#include <iostream>
using namespace std;
const int mxn = 1e5+5;
bool visited[mxn];
int sz[mxn];
vector<int> par;
vector<int> color, adj[mxn];
vector<pair<int, int>> new_edges;
void dfs(int u) {
visited[u] = true;
sz[u] = 1;
for(int v : adj[u]) if(!visited[v]) {
par[v] = u;
new_edges.push_back({u, v});
dfs(v); sz[u] += sz[v];
}
}
int find_centroid(int u, int n) {
visited[u] = true;
for(int v : adj[u]) if(!visited[v] && 2*sz[v]>n) {
return find_centroid(v, n);
}
return u;
}
void mark(int u, int m, vector<int> &where, int &cnt) {
if(cnt <= 0) return;
visited[u] = true;
where[u] = m;
cnt--;
for(int v : adj[u]) if(!visited[v]) {
mark(v, m, where, cnt);
}
}
int get_par(int u) {
return par[u] == u ? u : par[u] = get_par(par[u]);
}
bool join(int u, int v) {
u = get_par(u), v = get_par(v);
if(u == v) return false;
if(sz[u] < sz[v]) swap(u, v);
sz[u] += sz[v];
par[v] = u;
return true;
}
vector<int> find_split(int n, int a, int b, int c, vector<int> p, vector<int> q) {
for(int i = 0; i < p.size(); i++) {
int u = p[i], v = q[i];
adj[u].push_back(v), adj[v].push_back(u);
}
par.resize(n, -1);
dfs(0);
for(int i = 0; i < n; i++) {
adj[i].clear();
}
for(auto e : new_edges) {
adj[e.first].push_back(e.second);
adj[e.second].push_back(e.first);
}
fill(visited, visited+n, false);
int centroid = find_centroid(0, n);
vector<pair<int, int>> groups = {pair<int, int>{a, 1}, pair<int, int>{b, 2}, pair<int, int>{c, 3}};
sort(groups.begin(), groups.end());
color.resize(n, groups[2].second);
int small_group = -1;
int cnt;
if(par[centroid] != -1) sz[par[centroid]] = n-sz[centroid];
// continue adding non-tree edges until max sz exceeds min(a, b, c)
// or the edges run out
int mx = 0;
// current maximum
for(int u : adj[centroid]) mx = max(mx, sz[u]);
// initial parents (set sizes are equivalent to subtree ones)
fill(visited, visited+n, false);
for(int u : adj[centroid]) {
cnt = mxn;
mark(u, u, par, cnt);
}
// adding edges
for(int i = 0; i < p.size() && mx < groups[0].first; i++) {
if(p[i] != centroid && q[i] != centroid && join(p[i], q[i])) {
adj[p[i]].push_back(q[i]);
adj[q[i]].push_back(p[i]);
mx = max(mx, sz[get_par(p[i])]);
}
}
if(mx < groups[0].first) return vector<int>(n, 0);
fill(visited, visited+n, false);
for(int u : adj[centroid]) {
if(sz[u] >= groups[0].first) {
visited[centroid] = true;
cnt = groups[0].first;
mark(u, groups[0].second, color, cnt);
small_group = u;
break;
}
}
fill(visited, visited+n, false);
visited[small_group] = true;
cnt = groups[1].first;
mark(centroid, groups[1].second, color, cnt);
return color;
}
