#include "split.h"
#include <bits/stdc++.h>
using namespace std;
vector<vector<int>> new_edges;
int centroid_find(int n, vector<vector<int>> edges)
{
//create a spanning tree
//vector<vector<int>> new_edges(n);
vector<int> parent(n, -2);
vector<int> nodes;
new_edges = vector<vector<int>> (n);
nodes.push_back(0);
parent[0] = -1;
int curr = 0;
int total = 1;
while (curr < total)
{
int x = nodes[curr]; curr++;
for (int y: edges[x])
{
if (parent[y] == -2)
{
parent[y] = x;
nodes.push_back(y);
new_edges[x].push_back(y);
new_edges[y].push_back(x);
total++;
}
}
}
vector<int> subtreesize(n, 1);
for (int i=n-1; i>0; i--)
{
int x = nodes[i];
subtreesize[parent[x]] += subtreesize[x];
}
//now identify the centroid
int bes = n+1;
int cap = (n+1)/2;
int centroid = -1;
for (int i=0; i<n; i++)
if ((subtreesize[i]<bes) and (subtreesize[i]>=cap))
{
bes = subtreesize[i];
centroid = i;
}
return centroid;
}
vector<int> find_split_given_two_colors(int n, int a, int b, int c, vector<vector<int>> edges, vector<int> colors)
{
int col_4 = 0;
int col_5 = 0;
for (int i=0; i<n; i++)
{
if (colors[i]==4)
col_4++;
else
col_5++;
}
if (col_4 >= b)
{
for (int i=0; i<n; i++)
colors[i] = 9-colors[i];
swap(col_4, col_5);
}
int root, count;
root = 0;
count = 1;
while (colors[root]!=4)
root++;
colors[root] = 1;
queue<int> nodes4;
nodes4.push(root);
while (count < a)
{
int x = nodes4.front(); nodes4.pop();
for (int y: edges[x])
if ((colors[y] == 4) and (count < a))
{
colors[y] -= 3;
count++;
nodes4.push(y);
}
}
root = 0;
count = 1;
while (colors[root]!=5)
root++;
colors[root] = 2;
queue<int> nodes5;
nodes5.push(root);
while (count < b)
{
int x = nodes5.front(); nodes5.pop();
for (int y: edges[x])
if ((colors[y] == 5) and (count < b))
{
colors[y] -= 3;
count++;
nodes5.push(y);
}
}
for (int i=0; i<n; i++)
if (colors[i]>3)
colors[i] = 3;
return colors;
}
vector<int> find_split_wlog(int n, int a, int b, int c, vector<int> p, vector<int> q) {
vector<vector<int>> edges(n);
int e = p.size();
for (int i=0; i<e; i++)
{
int pp = p[i];
int qq = q[i];
edges[pp].push_back(qq);
edges[qq].push_back(pp);
}
int root = centroid_find(n, edges);
//cout << root << '\n';
//step 3: perform a dfs and identify whether its possible or not
vector<int> parent(n, -2);
stack<pair<int, int>> nodes;
vector<int> push_order;
vector<vector<int>> children(n);
parent[root] = -1;
nodes.push({root, 0});
push_order.push_back(root);
while (!nodes.empty())
{
pair<int, int> info = nodes.top();
int x = info.first; int y = info.second;
nodes.pop();
if ((y+1)<edges[x].size())
nodes.push({x, y+1});
int z = edges[x][y];
if (parent[z]==-2)
{
parent[z] = x;
nodes.push({z, 0});
push_order.push_back(z);
children[x].push_back(z);
}
}
vector<int> subtreesize(n, 1);
for (int i=n-1; i>0; i--)
{
int x = push_order[i];
subtreesize[parent[x]] += subtreesize[x];
//cout << x << subtreesize[x] << '\n';
}
//determine impossibility
int node_min = n+1;
int root_a = -1;
for (int i: children[root])
if ((i!=root) and (subtreesize[i] < node_min) and (subtreesize[i]>=a))
{
node_min = subtreesize[i];
root_a = i;
}
if (root_a == -1)
{
vector<int> ans(n, 0);
return ans;
}
//then a solution is possible
//do stuff with root_a
//cout << "reached here - 1\n";
parent = vector<int>(n, -2);
vector<int> color(n, -1);
vector<int> color_count(n, 0);
queue<int> nds;
parent[root] = -1;
nds.push(root);
while (!nds.empty())
{
int x = nds.front(); nds.pop();
for (int y: new_edges[x])
if (parent[y] == -2)
{
parent[y] = x;
int c = color[x];
if (x == root)
c = y;
color[y] = c;
color_count[c]++;
nds.push(y);
}
}
int big_alrdy = -1;
for (int i=0; i<n; i++)
if (color_count[i] >= a)
big_alrdy = i;
if (big_alrdy != -1)
{
//cout << "malaki na po ako\n";
vector<int> colorby2(n, 5);
for (int i=0; i<n; i++)
if (color[i] == big_alrdy)
colorby2[i] = 4;
return find_split_given_two_colors(n, a, b, c, edges, colorby2);
}
//cout << "reached here\n";
vector<vector<int>> color_edge(n);
for (int i=0; i<e; i++)
{
int ppp = color[p[i]];
int qqq = color[q[i]];
if ((ppp!=qqq) and (ppp >= 0) and (qqq >= 0))
{
color_edge[ppp].push_back(qqq);
color_edge[qqq].push_back(ppp);
}
}
int total = color_count[color[root_a]];
queue<int> colors_in;
vector<bool> colorby4(n, 0);
colors_in.push(color[root_a]);
colorby4[color[root_a]] = 1;
while (total < a)
{
int x = colors_in.front(); colors_in.pop();
for (int c: color_edge[x])
if ((colorby4[c]==0) and (total < a))
{
total += color_count[c];
colors_in.push(c);
colorby4[c] = 1;
}
}
vector<int> colorby2(n, 5);
for (int i=0; i<n; i++)
if (colorby4[color[i]])
colorby2[i] = 4;
return find_split_given_two_colors(n, a, b, c, edges, colorby2);
}
vector<int> find_split(int n, int a, int b, int c, vector<int> p, vector<int> q) {
//0 2 3 1
//112223
//223331
vector<int> rewrite = {0, 1, 2, 3};
for (int i=0; i<5; i++)
{
if (a>b)
{
swap(a, b);
swap(rewrite[1], rewrite[2]);
}
if (b>c)
{
swap(b, c);
swap(rewrite[2], rewrite[3]);
}
}
vector<int> ans = find_split_wlog(n, a, b, c, p, q);
for (int i=0; i<n; i++)
{
ans[i] = rewrite[ans[i]];
/*if (ans[i]==rewrite[0])
ans[i] = 0;
else if (ans[i]==rewrite[1])
ans[i] = 1;
else if (ans[i]==rewrite[2])
ans[i] = 2;
else if (ans[i]==rewrite[3])
ans[i] = 3;*/
}
return ans;
}
