#include "split.h"
#include <bits/stdc++.h>
using namespace std;
int numNodes;
vector<int> ans;
vector<int> edges1;
vector<int> edges2;
vector<vector<int>> neigh;
vector<int> par;
vector<int> stSize;
void initTree(int n, int p) {
par[n] = p;
stSize[n] = 1;
for (int i = 0; i < (int)neigh[n].size(); i++) {
if (neigh[n][i] != p) {
initTree(neigh[n][i], n);
stSize[n] += stSize[neigh[n][i]];
}
}
}
int remFill;
vector<bool> fillVis;
void fillDown(int n, int v) {
if (remFill == 0)
return;
remFill--;
ans[n] = v;
fillVis[n] = true;
for (int c : neigh[n]) {
if (c != par[n]) {
fillDown(c, v);
}
}
}
void fill(int n, int v) {
if (remFill == 0 || fillVis[n])
return;
fillVis[n] = true;
remFill--;
ans[n] = v;
for (int c : neigh[n]) {
fill(c, v);
}
}
bool solveTree(int a, int b, int c) {
stSize.resize(numNodes);
par.resize(numNodes);
initTree(0, -1);
int aRoot = -1;
int bRoot = -1;
for (int i = 0; i < numNodes; i++) {
int numUp = numNodes - stSize[i];
if (stSize[i] >= a && numUp >= b) {
aRoot = i;
break;
}
if (stSize[i] >= b && numUp >= a) {
bRoot = i;
break;
}
}
fill(ans.begin(), ans.end(), 2);
if (aRoot != -1) {
remFill = a;
fillDown(aRoot, 0);
remFill = b;
fill(par[aRoot], 1);
return true;
}
if (bRoot != -1) {
remFill = b;
fillDown(bRoot, 1);
remFill = a;
fill(par[bRoot], 0);
return true;
}
return false;
}
bool solveA1(int a, int b, int c) {
fill(ans.begin(), ans.end(), 2);
remFill = b;
fill(0, 1);
for (int& i : ans) {
if (i == 2) {
i = 0;
break;
}
}
return true;
}
vector<int> find_split(int _numNodes, int a, int b, int c, vector<int> p, vector<int> q) {
edges1 = move(p);
edges2 = move(q);
numNodes = _numNodes;
ans.resize(numNodes);
int numEdges = (int)edges1.size();
fillVis.resize(numNodes, false);
if (numEdges == numNodes && a != 1) {
edges1.pop_back();
edges2.pop_back();
numEdges--;
}
neigh.resize(numNodes);
for (int i = 0; i < (int)edges1.size(); i++) {
neigh[edges1[i]].push_back(edges2[i]);
neigh[edges2[i]].push_back(edges1[i]);
}
pair<int, int> abc[] = { { a, 1 }, { b, 2 }, { c, 3 } };
sort(abc, abc + 3);
bool can = false;
if (a == 1) {
can = solveA1(abc[0].first, abc[1].first, abc[2].first);
} else if (numEdges == numNodes - 1) {
can = solveTree(abc[0].first, abc[1].first, abc[2].first);
} else {
exit(1);
}
if (!can) {
fill(ans.begin(), ans.end(), 0);
} else {
for (int& i : ans) {
i = abc[i].second;
}
}
return ans;
}
