#include "split.h"
#include <vector>
#include <random>
#include <cassert>
#include <algorithm>
#include <bitset>
#include <vector>
using namespace std;
const int MAXN = 1e5 + 123;
vector<int> edg[MAXN];
int depth[MAXN];
int sz[MAXN];
int rt[MAXN];
int chld[MAXN];
bitset<MAXN> used;
void dfs(int v) {
used[v] = 1;
sz[v] = 1;
chld[v] = 0;
for (auto u : edg[v]) {
if (used[u])
continue;
++chld[v];
depth[u] = depth[v] + 1;
rt[u] = v;
dfs(u);
sz[v] += sz[u];
}
}
vector<int> split0(int n, int a, int b) {
for (int i = 0; i < n; ++i) {
if (rt[i] == i)
continue;
if (sz[i] >= b && n - sz[i] >= a) {
swap(a, b);
}
if (sz[i] >= a && n - sz[i] >= b) {
vector<int> ans(n, 0);
int ri = rt[i];
edg[i].erase(find(edg[i].begin(), edg[i].end(), ri));
edg[ri].erase(find(edg[ri].begin(), edg[ri].end(), i));
rt[i] = i;
rt[ri] = ri;
used = 0;
dfs(ri);
used = 0;
dfs(i);
--chld[i];
--chld[ri];
vector<int> leafa;
vector<int> leafb;
for (int j = 0; j < n; ++j) {
if (used[j]) {
ans[j] = 0;
if (chld[j] == 0)
leafa.push_back(j);
continue;
}
ans[j] = 1;
if (chld[j] == 0) {
leafb.push_back(j);
}
}
int c = n - a - b;
int sza = sz[i];
int szb = sz[ri];
while (c > 0 && sza > a && leafa.size()) {
int v = leafa.back();
leafa.pop_back();
--c;
--sza;
ans[v] = 2;
for (auto u : edg[v]) {
if (--chld[u] == 0)
leafa.push_back(u);
}
}
while (c > 0 && szb > b && leafb.size()) {
int v = leafb.back();
leafb.pop_back();
--c;
--szb;
ans[v] = 2;
for (auto u : edg[v]) {
if (--chld[u] == 0)
leafb.push_back(u);
}
}
return ans;
}
}
return {};
}
vector<int> solve(int n, int a, int b, int c, int root) {
used = 0;
depth[root] = 0;
rt[root] = root;
dfs(root);
vector<int> vls = {a, b, c};
vector<int> rpl;
if (c >= a && c >= b)
rpl = {0, 1, 2};
else if (b >= a && b >= c)
rpl = {0, 2, 1};
else
rpl = {1, 2, 0};
if (vls[rpl[0]] > vls[rpl[1]])
swap(rpl[0], rpl[1]);
for (int i = 0; i < n; ++i)
edg[i].clear();
for (int i = 0; i < n; ++i) {
if (rt[i] == i)
continue;
edg[rt[i]].push_back(i);
edg[i].push_back(rt[i]);
}
//assert(vls[rpl[0]] <= vls[rpl[1]] && vls[rpl[1]] <= vls[rpl[2]]);
vector<int> res = split0(n, vls[rpl[0]], vls[rpl[1]]);
if (res.empty()) {
return {};
} else {
vector<int> cnts(3, 0);
for (int i = 0; i < n; ++i)
++cnts[res[i]];
if (cnts[1] < cnts[0]) {
for (int i = 0; i < n; ++i) {
if (res[i] < 2)
res[i] = 1 ^ res[i];
}
}
vector<int> rrpl(3);
for (int i = 0; i < 3; ++i)
rrpl[rpl[i]] = i;
for (int i = 0; i < n; ++i)
res[i] = rpl[res[i]] + 1;
return res;
}
return res;
}
mt19937 rng(177013);
vector<int> find_split(int n, int a, int b, int c, vector<int> p, vector<int> q) {
int t = 30000000 / (n + p.size());
while (t--) {
for (int i = 0; i < n; ++i)
edg[i].clear();
for (int i = 0; i < p.size(); ++i) {
edg[p[i]].push_back(q[i]);
edg[q[i]].push_back(p[i]);
}
for (int i = 0; i < n; ++i)
shuffle(edg[i].begin(), edg[i].end(), rng);
int root = rng() % n;
vector<int> res = solve(n, a, b, c, root);
if (res.size())
return res;
}
}
