#include <bits/stdc++.h>
using namespace std;
const int n_max = 2e5+1e3;
int visit[n_max];
int low[n_max];
int subtree[n_max];
vector<int> adj[n_max];
vector<int> child[n_max];
int parent[n_max];
int time_global = 0;
void dfs(int x) {
subtree[x] = 1;
visit[x] = time_global;
time_global++;
low[x] = x;
for(int i: adj[x]) {
if(visit[i]==-1) {
parent[i] = x;
child[x].push_back(i);
dfs(i);
subtree[x] += subtree[i];
low[x] = min(low[x], low[i]);
} else if (i != parent[x]) {
low[x] = min(low[x], visit[i]);
}
}
}
vector<int> find_split(int n, int a, int b, int c, vector<int> p, vector<int> q) {
// assume a<=b<=c
vector<pair<int,int>> vp = {make_pair(a,1),make_pair(b,2),make_pair(c,3)};
sort(vp.begin(), vp.end());
int small = vp[0].second;
int mid = vp[1].second;
int big = vp[2].second;
a = vp[0].first;
b = vp[1].first;
c = vp[2].first;
memset(visit, -1, sizeof(visit));
for(int i=0; i<(int) p.size(); i++) {
adj[p[i]].push_back(q[i]);
adj[q[i]].push_back(p[i]);
}
dfs(0);
vector<int> ans(n,0);
for(int i=0; i<n; i++) {
// assume c>=a
// check: a+c or b+c?
if(subtree[i]>=a && subtree[i]<=b+c) {
if(subtree[i]>a+c) {
swap(a,b);
swap(small,mid);
}
queue<int> bfs;
bfs.push(i);
int ctr = 0;
while(!bfs.empty()) {
int next = bfs.front();
ans[next] = (ctr<a ? small : big);
ctr++;
bfs.pop();
for(int k: child[next]) {
bfs.push(k);
}
}
bfs.push(0);
ctr = 0;
while(!bfs.empty()) {
int next = bfs.front();
ans[next] = (ctr<b ? mid : big);
ctr++;
bfs.pop();
for(int k: child[next]) {
if(ans[k]==0) {
bfs.push(k);
}
}
}
for(int k=0; k<n; k++) {
assert(ans[k]!=0);
}
return ans;
} else if(subtree[i]>b+c) {
// if have heavy node with all light children
bool all_light = true;
for(int j: child[i]) {
if(subtree[j]>=a) {
all_light = false;
// break;
}
}
if(all_light) {
if(i==0) {
assert(child[i].size()>=3);
return ans;
}
int minsize = 1;
vector<int> tmp; // consists of all the children to be taken
for(int j: child[i]) {
// add all those with bad lowtime
if(low[j]>= visit[i]) {
minsize += subtree[j];
tmp.push_back(j);
}
}
if(minsize > b+c) {
return ans;
}
for(int j: child[i]) {
// add the relevant subtrees
if(low[j]<visit[i] && minsize < b) {
minsize += subtree[j];
tmp.push_back(j);
}
}
assert(minsize >= b);
assert(minsize <= b+c);
// now for each subtree, color its children completely except that some are colored with 3
int ctr = 1;
ans[i] = mid;
queue<int> bfs;
for(int j: tmp) {
bfs.push(j);
while(!bfs.empty()) {
int next = bfs.front();
assert(ans[next]==0);
ans[next] = (ctr<b ? mid : big);
ctr++;
bfs.pop();
for(int k: child[next]) {
bfs.push(k);
}
}
}
assert(ctr==b);
ctr = 1;
bfs.push(0); // we know that 0 is not the root
ans[0] = small;
while(!bfs.empty()) {
int next = bfs.front();
bfs.pop();
for(int k: adj[next]) {
if(ans[k]==0) {
bfs.push(k);
ans[k] = (ctr<a ? small : big);
ctr++;
}
}
}
assert(ctr==a);
for(int k=0; k<n; k++) {
//assert(ans[k]!=0);
}
return ans;
}
}
}
assert(false);
return p;
}
int main()
{
/*int n=999;
int a=((rand()%(n-3))/2)+1;
int b=((rand()%(n-3))/2)+1;
int c=n-a-b;
assert(a>=1);
assert(b>=1);
assert(c>=1);
assert(a+b+c==n);
auto x = gen_graph(n);
vector<int> p,q;
p = x.first;
q = x.second;*/
int n = 10;
int a=3;
int b=3;
int c=4;
/*vector<int> p = {0,2,1,3,1,4,1,1,6,1,7,1,8,1,9};
vector<int> q = {1,0,3,0,4,0,5,6,0,7,0,8,0,9,0};*/
vector<int> p = {0,1,1,1,1,1,1,1,1,7,6};
vector<int> q = {1,2,3,4,5,6,7,8,9,0,0};
vector<int> v = find_split(n,a,b,c,p,q);
//verify_ans(n,a,b,c,p,q,v);
for(int i: v) {
cout << i << " ";
}
return 0;
}
Compilation message
/tmp/ccStakG4.o: In function `main':
grader.cpp:(.text.startup+0x0): multiple definition of `main'
/tmp/ccktPLNN.o:split.cpp:(.text.startup+0x0): first defined here
collect2: error: ld returned 1 exit status
