#include <bits/stdc++.h>
#include "split.h"
using namespace std;
const int MAXN = 100010;
int subtree_size[MAXN];
vector<int> adj[MAXN];
int parent[MAXN];
void dfs(int src = 0, int prev = -1) {
//printf("dfs %d %d\n", src, prev);
subtree_size[src] = 1;
parent[src] = prev;
for(auto &xt : adj[src]) {
if(xt - prev) {
dfs(xt, src);
subtree_size[src] += subtree_size[xt];
}
}
}
vector<int> get_solution(vector<int> v1, int n1, vector<int> v2, int n2, int n) {
vector<int> ans(n, 6 - n1 - n2);
for(auto &x : v1) {
ans[x] = n1;
}
for(auto &x : v2) {
ans[x] = n2;
}
return ans;
}
bool marked[MAXN];
vector<int> extract(int src, int wrong, int sz) {
marked[wrong] = true;
vector<int> ans;
stack<int> nodes;
nodes.push(src);
marked[src] = true;
while((int)ans.size() < sz) {
int node = nodes.top();
nodes.pop();
ans.push_back(node);
for(auto &x : adj[node]) {
if(!marked[x]) {
marked[x] = true;
nodes.push(x);
}
}
}
for(auto &x : ans) {
marked[x] = false;
}
marked[wrong] = false;
return ans;
}
vector<int> blank(int n) {
vector<int> ans(n, 0);
return ans;
}
int dsu[MAXN];
void init_set(int x) {
dsu[x] = x;
}
int get(int x) {
if(dsu[x] == x) return x;
return dsu[x] = get(dsu[x]);
}
bool unite(int x, int y) {
int xr = get(x), yr = get(y);
if(xr == yr) return false;
if(rand() % 2) swap(xr, yr);
dsu[xr] = yr;
return true;
}
void create_spanning_tree(int n, vector<int> &st_p, vector<int> &st_q, vector<int> p, vector<int> q) {
for(int i = 0; i < n; i++) {
init_set(i);
}
int m = (int)p.size();
for(int i = 0; i < m; i++) {
if(unite(p[i], q[i])) {
st_p.push_back(p[i]);
st_q.push_back(q[i]);
}
}
assert((int)st_p.size() == n - 1);
}
int belongs[MAXN];
set<int> comp_adj[MAXN];
int cnt_nodes[MAXN];
int comp_size;
vector<int> comp_comp;
bool comp_mark[MAXN];
void dfs2(int src) {
comp_size += cnt_nodes[src];
comp_mark[src] = true;
comp_comp.push_back(src);
for(auto &xt : comp_adj[src]) {
if(!comp_mark[xt]) {
dfs2(xt);
}
}
}
vector<int> find_split(int n, int a, int b, int c, vector<int> p, vector<int> q) {
pair<int, int> arr[] = {{a, 1}, {b, 2}, {c, 3}};
sort(arr, arr + 3);
int A = arr[0].first, B = arr[1].first, C = arr[2].first;
vector<int> st_p, st_q;
create_spanning_tree(n, st_p, st_q, p, q);
int m = (int)p.size();
for(int i = 0; i < n - 1; i++) {
adj[st_p[i]].push_back(st_q[i]);
adj[st_q[i]].push_back(st_p[i]);
}
dfs();
//for(int i = 0; i < n; i++) {
// printf("subtree_size[%d] = %d\n", i, subtree_size[i]);
//}
for(int i = 0; i < m; i++) {
int u = p[i], v = q[i];
if(subtree_size[u] > subtree_size[v]) swap(u, v);
int under = subtree_size[u];
int rest = n - under;
if(under >= A && rest >= A) {
//printf("found for edge %d - %d\n", u, v);
//we've got a solution
vector<int> set_A, set_B;
if(under < rest) {
set_A = extract(u, v, A);
set_B = extract(v, u, B);
} else {
set_A = extract(v, u, A);
set_B = extract(u, v, B);
}
return get_solution(set_A, arr[0].second, set_B, arr[1].second, n);
}
}
//there is no edge connecting two subtrees of sizes >= a
//meaning that it's a star with all subtrees attached of size < a
//in tree, there is no such solution
for(int node = 0; node < n; node++) {
vector<pair<int, int>> sizes;
int sm = 0;
for(auto &xt : adj[node]) {
if(xt != parent[node]) {
sizes.push_back({subtree_size[xt], xt});
sm += subtree_size[xt];
}
}
if(node) {
sizes.push_back({n - 1 - sm, parent[node]});
}
int mx = 0;
for(auto &pr : sizes) {
mx = max(mx, pr.first);
}
if(mx < A) {
printf("I am deciding on node %d\n", node);
for(int i = 0; i < (int)sizes.size(); i++) {
vector<int> comp_here = extract(sizes[i].second, node, sizes[i].first);
for(auto &x : comp_here) {
belongs[x] = i + 1;
marked[x] = false;
}
marked[node] = false;
cnt_nodes[i + 1] = sizes[i].first;
}
int new_nodes = (int)sizes.size();
for(int i = 0; i < m; i++) {
if(p[i] != node && q[i] != node) {
comp_adj[belongs[p[i]]].insert(belongs[q[i]]);
comp_adj[belongs[q[i]]].insert(belongs[p[i]]);
}
}
for(int i = 1; i <= new_nodes; i++) {
if(marked[i]) continue;
comp_size = 0;
comp_comp.clear();
dfs2(i);
for(auto &x : comp_comp) {
comp_mark[x] = false;
}
if(comp_size >= A) {
//we have a solution
int start_node = sizes[i - 1].second;
vector<int> ans_A = extract(start_node, node, A);
for(auto &x : ans_A) {
marked[x] = true;
}
vector<int> ans_B = extract(node, MAXN - 1, B);
return get_solution(ans_A, arr[0].second, ans_B, arr[1].second, n);
} else {
return blank(n);
}
}
}
}
assert(false);
return blank(n);
}
