#include <bits/stdc++.h>
using namespace std;
//macros
typedef long long ll;
typedef pair<int, int> ii;
typedef tuple<int, int, int> iii;
typedef vector<int> vi;
typedef vector<ii> vii;
typedef vector<iii> viii;
typedef vector<ll> vll;
#define REP(a,b,c) for(int a=int(b); a<int(c); a++)
#define RE(a,c) REP(a,0,c)
#define RE1(a,c) REP(a,1,c+1)
#define REI(a,b,c) REP(a,b,c+1)
#define REV(a,b,c) for(int a=int(c-1); a>=int(b); a--)
#define INF 1e9
#define pb push_back
#define fi first
#define se second
mt19937 rng(chrono::steady_clock::now().time_since_epoch().count());
const int MX = 4e5;
int n, m, a[3], SA[3], r;
vi adj[MX], chl[MX], tAdj[MX];
bitset<MX> visited;
int sz[MX];
int cent;
int subN[MX];
set<int> subAdj[MX];
set<int> curSub;
int curSubSz;
vi ans;
void dfsChl(int u=r) {
visited[u] = 1;
for(int v:adj[u]) {
if(visited[v]) continue;
dfsChl(v);
chl[u].pb(v);
tAdj[u].pb(v);
tAdj[v].pb(u);
}
}
void dfsSize(int u=r) {
sz[u] = 1;
for(int v:chl[u]) {
dfsSize(v);
sz[u] += sz[v];
}
}
int getCentroid(int u=r) {
for(int v:chl[u])
if(sz[v] > n/2)
return getCentroid(v);
return u;
}
void dfsChl2(int u) {
visited[u] = 1;
for(int v:tAdj[u]) if(!visited[v]) dfsChl2(v), chl[u].pb(v);
}
void dfsSubNumb(int u, int s) {
subN[u] = s;
for(int v:chl[u]) dfsSubNumb(v, s);
}
void dfsSubSize(int u) {
if(visited[u]) return;
visited[u] = 1;
if(curSubSz >= a[SA[0]]) return;
curSubSz += sz[u];
curSub.insert(u);
for(int v:subAdj[u]) {
dfsSubSize(v);
if(curSubSz >= a[SA[0]]) return;
}
}
void dfsAns1(int u) {
if(!curSub.count(subN[u])) return;
if(visited[u]) return;
if(a[SA[0]]==0) return;
visited[u] = 1;
a[SA[0]]--;
ans[u] = 0;
for(int v:adj[u]) dfsAns1(v);
}
void dfsAns2(int u) {
if(curSub.count(subN[u])) return;
if(visited[u]) return;
if(a[SA[1]]==0) return;
visited[u] = 1;
a[SA[1]]--;
ans[u] = 1;
for(int v:adj[u]) dfsAns2(v);
}
vi find_split(int N, int A, int B, int C, vi p, vi q) {
RE(i,MX) subN[i]=0, sz[i]=0;
curSubSz=0;
visited.reset();
r=rng()%N;
n=N; m=p.size();
a[0] = A;
a[1] = B;
a[2] = C;
RE(i,3) SA[i]=i;
sort(SA,SA+3,[](int i, int j) {return a[i]<a[j];});
RE(i,m) {
adj[p[i]].pb(q[i]);
adj[q[i]].pb(p[i]);
}
visited.reset();
dfsChl();
dfsSize();
cent = getCentroid();
RE(i,n) chl[i].clear();
visited.reset();
dfsChl2(cent);
dfsSize(cent);
for(int v:chl[cent]) dfsSubNumb(v, v);
RE(u,n) {
if(u == cent) continue;
for(int v:adj[u]) if(subN[u] != subN[v]) {
if(v == cent) continue;
subAdj[subN[u]].insert(subN[v]);
subAdj[subN[v]].insert(subN[u]);
}
}
visited.reset();
for(int v:chl[cent]) {
if(visited[v]) continue;
curSubSz = 0;
curSub.clear();
dfsSubSize(v);
if(curSubSz >= a[SA[0]]) break;
curSub.clear();
}
ans.resize(n);
if(curSub.empty()) {
RE(i,n) ans[i] = 0;
} else {
RE(i,n) ans[i] = 2;
visited.reset();
dfsAns1(*curSub.begin());
dfsAns2(cent);
RE(i,n) ans[i] = SA[ans[i]]+1;
}
return ans;
}
int main() {
vi ans = find_split(9, 4, 2, 3, {0, 0, 0, 0, 0, 0, 1, 3, 4, 5}, {1, 2, 3, 4, 6, 8, 7, 7, 5, 6});
bool f=1;
for(int v:ans) cout<<(f?"":" ")<<v, f=0; cout<<endl;
}
Compilation message
split.cpp: In function 'int main()':
split.cpp:152:2: warning: this 'for' clause does not guard... [-Wmisleading-indentation]
for(int v:ans) cout<<(f?"":" ")<<v, f=0; cout<<endl;
^~~
split.cpp:152:43: note: ...this statement, but the latter is misleadingly indented as if it were guarded by the 'for'
for(int v:ans) cout<<(f?"":" ")<<v, f=0; cout<<endl;
^~~~
/tmp/ccUPUfns.o: In function `main':
grader.cpp:(.text.startup+0x0): multiple definition of `main'
/tmp/cciktvzp.o:split.cpp:(.text.startup+0x0): first defined here
collect2: error: ld returned 1 exit status
