#include "split.h"
#include <vector>
#include <algorithm>
#include <iostream>
#include <queue>
using namespace std;
/*
1. Find an arbitrary spanning tree.
2. Find centroid.
3. Find subtree sizes.
4. Check if any subtree has size >= A. If yes, return answer.
5. (2A + B <= N) Using remaining edges, join subtrees.
If any set of joined subtrees has size >= A, return answer.
6. Return NO.
*/
int A, A_index, B, B_index, C, C_index;
struct group
{
int v;
int i;
};
bool operator < (group A, group B)
{
return A.v < B.v;
}
const int maxN = 100'000;
vector<int> edge[1+maxN];
int N;
vector<int> st_edge[1+maxN];
vector<int> visit(1+maxN, 0);
vector<int> temp_size(1+maxN, 1);
void dfs(int u)
{
visit[u] = 1;
for(int v: edge[u])
{
if(visit[v]) continue;
st_edge[u].push_back(v);
st_edge[v].push_back(u);
dfs(v);
temp_size[u] += temp_size[v];
}
}
vector<int> new_size(1+maxN, 1);
vector<int> centroid_depth(1+maxN);
void new_dfs(int u)
{
visit[u] = 1;
for(int v: st_edge[u])
{
if(visit[v]) continue;
centroid_depth[v] = centroid_depth[u] + 1;
new_dfs(v);
new_size[u] += new_size[v];
}
}
int centroid;
vector<int> res;
int A_count = 0, B_count = 0;
void A_dfs(int u)
{
if(A_count == A) return;
res[u] = A_index;
A_count++;
for(int v: st_edge[u])
{
if(v == centroid) continue;
if(res[v] == A_index) continue;
A_dfs(v);
}
}
void B_dfs(int u)
{
if(B_count == B) return;
res[u] = B_index;
B_count++;
for(int v: st_edge[u])
{
if(res[v] != C_index) continue;
B_dfs(v);
}
}
int subtree_ct = 0;
vector<int> subtree_index(1+maxN, 0);
vector<int> subtree_size(1);
void final_tree_dfs(int u)
{
subtree_index[u] = subtree_ct;
subtree_size[subtree_ct]++;
for(int v: st_edge[u])
{
if(v == centroid || subtree_index[v]) continue;
final_tree_dfs(v);
}
}
vector<int> extra_edge[1+maxN];
vector<int> find_split(int n, int a, int b, int c, vector<int> p, vector<int> q)
{
//PART 0: INPUT
N = n;
int M = p.size();
for(int j = 0; j < M; j++)
{
edge[ p[j] ].push_back( q[j] );
edge[ q[j] ].push_back( p[j] );
}
vector<group> G{group{a, 1}, group{b, 2}, group{c, 3}};
sort(G.begin(), G.end());
A = G[0].v;
A_index = G[0].i;
B = G[1].v;
B_index = G[1].i;
C = G[2].v;
C_index = G[2].i;
//PART 1: FIND SPANNING TREE.
dfs(0);
// cerr << "check A\n";
//PART 2: FIND CENTROID.
for(centroid = 0; centroid < N; centroid++)
{
bool good = 1;
if(2 * (N - temp_size[centroid]) > N) good = 0;
for(int v: st_edge[centroid])
{
if(temp_size[v] > temp_size[centroid]) continue;
if(2 * temp_size[v] > N) good = 0;
}
if(good) break;
}
// cerr << "check\n";
// cerr << C_index << '\n';
//PART 3: FIND SUBTREE SIZES
visit = vector<int>(N, 0);
centroid_depth[centroid] = 0;
new_dfs(centroid);
//PART 4: CHECK IF ANY SUBTREE HAS SIZE >= A, IF YES RETURN
res = vector<int>(N, C_index);
for(int X: st_edge[centroid])
{
if(A <= new_size[X] && B <= N - new_size[X])
{
A_dfs(X);
B_dfs(centroid);
return res;
}
}
//PART 5: CHECK IF ANY COMBINATION OF SUBTREES HAVE SIZE >= A.
for(int u: st_edge[centroid])
{
subtree_ct++;
subtree_size.push_back(0);
final_tree_dfs(u);
}
for(int j = 0; j < M; j++)
{
if(subtree_index[ p[j] ] == subtree_index[ q[j] ]) continue;
if(p[j] == centroid || q[j] == centroid) continue;
extra_edge[ subtree_index[ p[j] ] ].push_back( subtree_index[ q[j] ] );
extra_edge[ subtree_index[ q[j] ] ].push_back( subtree_index[ p[j] ] );
}
vector<int> group_visit(1+maxN, 0);
vector<int> group_size(1+maxN, 0);
int curr = 0;
queue<int> tbv;
bool solution_exists = 0;
vector<int> original_A(1+maxN, 0);
for(int u = 1; u <= subtree_ct; u++)
{
if(group_visit[u]) continue;
curr++;
tbv.push(u);
while(!tbv.empty())
{
int U = tbv.front();
group_visit[U] = curr;
group_size[curr] += subtree_size[U];
if(group_size[curr] >= A)
{
solution_exists = 1;
for(int i = 0; i < N; i++)
{
if(group_visit[ subtree_index[i] ] == curr)
{
original_A[i] = 1;
}
}
break;
}
if(solution_exists) break;
for(int v: extra_edge[U])
{
if(group_visit[v]) continue;
tbv.push(v);
}
}
}
if(!solution_exists) return vector<int>(N, 0);
res = vector<int>(N, C_index);
int count_of_A = 0, count_of_B = 0;
while(!tbv.empty()) tbv.pop();
for(int x: st_edge[centroid])
{
if(original_A[x])
{
tbv.push(x);
break;
}
}
while(!tbv.empty())
{
int u = tbv.front();
res[u] = A_index;
count_of_A++;
if(count_of_A == A) break;
for(int v: edge[u])
{
if(original_A[v])
{
tbv.push(v);
}
}
}
while(!tbv.empty()) tbv.pop();
tbv.push(centroid);
while(!tbv.empty())
{
int u = tbv.front();
res[u] = B_index;
count_of_B++;
if(count_of_B == B) break;
for(int v: edge[u])
{
if(res[v] == C_index)
{
tbv.push(v);
}
}
}
return res;
}
