#include <bits/stdc++.h>
using namespace std;
typedef pair<int, int> pii;
vector<int> edges[100005];
int ufpar[100005];
int ufsz[100005];
int find(int x) {
return ufpar[x] == x ? x : (ufpar[x] = find(ufpar[x]));
}
bool merge(int x, int y) {
x = find(x);
y = find(y);
if(x == y) return false;
if(ufsz[x] < ufsz[y]) swap(x, y);
ufsz[x] += ufsz[y];
ufpar[y] = x;
return true;
}
// spanning tree is rooted at 0
int treeSize[100005];
int dfsParent[100005];
bool dfsSeen[100005];
int assignedColor[100005];
int ret[100005];
bool ufcolor(int curr, int need, int want) {
ret[curr] = want;
if(--need == 0) return true;
queue<int> q;
q.push(curr);
while(!q.empty()) {
curr = q.front(); q.pop();
for(int out: edges[curr]) {
if(find(out) != find(curr)) continue;
if(ret[out]) continue;
ret[out] = want;
if(--need == 0) return true;
q.push(out);
}
}
return false;
}
vector<int> consolidate(int n) {
vector<int> ans;
for(int i = 0; i < n; i++) {
assert(ret[i] >= 1 && ret[i] <= 3);
ans.push_back(ret[i]);
}
return ans;
}
void downColor(int curr, int par, int& need, int want) {
ret[curr] = want;
assert(need > 0);
if(--need == 0) return;
for(int out: edges[curr]) {
if(out == par || need == 0) continue;
downColor(out, curr, need, want);
}
}
bool upColor(int curr, int need, int want) {
queue<int> q;
q.push(curr);
ret[curr] = want;
if(--need == 0) return true;
while(!q.empty()) {
curr = q.front(); q.pop();
for(int out: edges[curr]) {
if(ret[out]) continue;
ret[out] = want;
if(--need == 0) return true;
q.push(out);
}
}
return false;
}
void fillRest(int n, int want) {
for(int i = 0; i < n; i++) {
if(ret[i] == 0) ret[i] = want;
}
}
void dfs(int curr, int par) {
treeSize[curr] = 1;
dfsSeen[curr] = true;
dfsParent[curr] = par;
for(int out: edges[curr]) {
if(dfsSeen[out]) continue;
dfs(out, curr);
treeSize[curr] += treeSize[out];
}
}
vector<int> find_split(int n, int a, int b, int c, vector<int> p, vector<int> q) {
for(int i = 0; i < p.size(); i++) {
edges[p[i]].push_back(q[i]);
edges[q[i]].push_back(p[i]);
}
dfs(0, -1);
vector<pii> sizes;
sizes.emplace_back(a, 1);
sizes.emplace_back(b, 2);
sizes.emplace_back(c, 3);
sort(sizes.begin(), sizes.end());
// is there a subtree such that it can be colored with size a and the other half can be colored with size b
for(int i = 1; i < n; i++) {
if(treeSize[i] >= sizes[0].first && n-treeSize[i] >= sizes[1].first) {
downColor(i, dfsParent[i], sizes[0].first, sizes[0].second);
assert(upColor(dfsParent[i], sizes[1].first, sizes[1].second));
fillRest(n, sizes[2].second);
return consolidate(n);
}
}
// find a node such that, when deleted, all trees have size < a
int critical = -1;
for(int i = 0; i < n; i++) {
int largestTree = 0;
int parentSize = n-1;
for(int out: edges[i]) {
if(dfsParent[out] != i) continue;
largestTree = max(largestTree, treeSize[out]);
parentSize -= treeSize[out];
}
largestTree = max(largestTree, parentSize);
if(largestTree >= sizes[0].first) continue;
critical = i;
break;
}
assert(critical != -1);
for(int i = 0; i < n; i++) {
ufpar[i] = i;
ufsz[i] = 1;
}
for(int i = 1; i < n; i++) {
if(i == critical || dfsParent[i] == critical) continue;
assert(merge(i, dfsParent[i]));
assert(ufsz[find(i)] < sizes[0].first);
}
for(int i = 0; i < n; i++) {
for(int out: edges[i]) {
if(i == critical || out == critical) continue;
if(!merge(i, out)) continue;y
if(ufsz[find(i)] >= sizes[0].first) {
ufcolor(i, sizes[0].first, sizes[0].second);
assert(upColor(critical, sizes[1].first, sizes[1].second));
fillRest(n, sizes[2].second);
return consolidate(n);
}
}
}
vector<int> ans;
ans.resize(n);
return ans;
}
Compilation message
split.cpp: In function 'std::vector<int> find_split(int, int, int, int, std::vector<int>, std::vector<int>)':
split.cpp:104:20: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
for(int i = 0; i < p.size(); i++) {
~~^~~~~~~~~~
split.cpp:151:7: warning: this 'if' clause does not guard... [-Wmisleading-indentation]
if(!merge(i, out)) continue;y
^~
split.cpp:151:35: note: ...this statement, but the latter is misleadingly indented as if it were guarded by the 'if'
if(!merge(i, out)) continue;y
^
split.cpp:151:35: error: 'y' was not declared in this scope
