This submission is migrated from previous version of oj.uz, which used different machine for grading. This submission may have different result if resubmitted.
#include <bits/stdc++.h>
#include "split.h"
using namespace std;
using pii = pair<int, int>;
vector<int> tin, low, root, val;
int timer;
vector<vector<pii> > g;
vector<vector<int> > cg, components(1);
int comp_cnt = 0;
vector<bool> vis;
vector<bool> is_bridge;
vector<int> sub;
pair<int, int> sz[3];
int cnt[3];
int ansv = -1, ansp = -1;
int sub_dfs(int v, int p) {
sub[v] = val[v];
for(int to : cg[v]) {
if(to != p) {
sub[v] += sub_dfs(to, v);
}
}
return sub[v];
}
void finder_dfs(int v, int p) {
if(ansv != -1) {
return;
}
const int subtree = sub[v]-val[v], uptree = sub[0] - sub[v];
const int diff = max(0, sz[0].first - min(subtree, uptree)) + max(0, sz[1].first - max(subtree, uptree));
if(diff <= val[v]) {
ansv = v;
ansp = p;
return;
}
for(int to : cg[v]) {
if(to != p) {
finder_dfs(to, v);
if(ansv != -1) {
return;
}
}
}
}
void bridge_dfs(int v, int p) {
assert(!vis[v]);
vis[v] = true;
tin[v] = low[v] = ++timer;
for(auto e : g[v]) {
int to = e.first, idx = e.second;
if(to == p) {
continue;
}
if(!vis[to]) {
bridge_dfs(to, v);
if(low[to] > tin[v]) {
// cout << v << ' ' << to << endl;
is_bridge[idx] = true;
}
low[v] = min(low[v], low[to]);
}
else {
low[v] = min(low[v], low[to]);
}
}
}
void comp_dfs(int v, int cur_comp) {
assert(!vis[v]);
vis[v] = true;
root[v] = cur_comp;
components[cur_comp].push_back(v);
val[cur_comp]++;
for(const auto& e : g[v]) {
int to = e.first, idx = e.second;
if(vis[to]) {
continue;
}
if(is_bridge[idx]) {
comp_cnt++;
components.emplace_back();
comp_dfs(to, comp_cnt);
}
else {
comp_dfs(to, cur_comp);
}
}
}
void finalize_dfs(vector<int>& res, int v, int p, int group) {
if(cnt[group] >= sz[group].first) {
assert(cnt[group] <= sz[group].first);
return;
}
assert(!res[v]);
assert(0 <= group and group < 3);
res[v] = sz[group].second;
cnt[group]++;
for(const auto& e : g[v]) {
int to = e.first, idx = e.second;
if(root[to] == p or res[to] != 0) {
continue;
}
finalize_dfs(res, to, v, group);
if(cnt[group] >= sz[group].first) {
assert(cnt[group] <= sz[group].first);
return;
}
}
}
void bounded_dfs(vector<int>& res, int v, int group) {
if(cnt[group] >= sz[group].first) {
assert(cnt[group] <= sz[group].first);
return;
}
assert(!res[v]);
assert(0 <= group and group < 3);
res[v] = sz[group].second;
cnt[group]++;
for(const auto& e : g[v]) {
int to = e.first, idx = e.second;
if(root[to] != root[v] or res[to] != 0) {
continue;
}
finalize_dfs(res, to, v, group);
if(cnt[group] >= sz[group].first) {
assert(cnt[group] <= sz[group].first);
return;
}
}
}
vector<signed> find_split(signed n, signed a, signed b, signed c, vector<signed> p, vector<signed> q) {
sz[0] = {a, 1};
sz[1] = {b, 2};
sz[2] = {c, 3};
sort(sz, sz+3);
vector<signed> res(n);
int m = p.size();
sub.assign(n, 0);
val.assign(n, 0);
g.assign(n, vector<pii>());
vis.assign(n, false);
tin.assign(n, -1);
low.assign(n, -1);
root.assign(n, -1);
is_bridge.assign(m, false);
for(int i = 0; i < m; ++i) {
const int u = p[i];
const int v = q[i];
g[u].push_back({v, i});
g[v].push_back({u, i});
}
bridge_dfs(0, -1);
vis.assign(n, false);
comp_dfs(0, comp_cnt);
cg.assign(comp_cnt + 3, vector<int>());
for(int u = 0; u < n; ++u) {
for(const auto& e : g[u]) {
int v = e.first;
assert(root[u] != -1 and root[v] != -1);
if(root[u] == root[v]) {
continue;
}
cg[root[u]].push_back(root[v]);
}
}
sub_dfs(0, -1);
finder_dfs(0, -1);
// cout << ansv << ' ' << ansp << endl;
if(ansv == -1) {
// cout << "Hi";
assert(ansp == -1);
return res;
}
const int subtree = sub[ansv]-val[ansv], uptree = sub[0] - sub[ansv];
if(subtree < uptree) {
for(int child : cg[ansv]) {
if(child == ansp) {
continue;
}
finalize_dfs(res, components[child].back(), ansv, 0);
}
if(ansp != -1) {
finalize_dfs(res, ansp, ansv, 1);
}
}
else {
for(int child : cg[ansv]) {
if(child == ansp) {
continue;
}
finalize_dfs(res, components[child].back(), ansv, 1);
}
if(ansp != -1) {
finalize_dfs(res, ansp, ansv, 0);
}
}
for(int v : components[ansv]) {
if(!res[v]) {
bounded_dfs(res, v, 0);
break;
}
}
for(int v : components[ansv]) {
if(!res[v]) {
bounded_dfs(res, v, 1);
break;
}
}
for(int& x : res)
if(!x)
x = sz[2].second;
return res;
}
/*
6 5
1 3 2
0 1
0 2
0 3
0 4
0 5
9 10
3 3 3
0 1
0 2
0 3
0 4
0 6
0 8
1 7
3 7
4 5
5 6
*/
Compilation message (stderr)
split.cpp: In function 'void finalize_dfs(std::vector<int>&, int, int, int)':
split.cpp:108:21: warning: unused variable 'idx' [-Wunused-variable]
108 | int to = e.first, idx = e.second;
| ^~~
split.cpp: In function 'void bounded_dfs(std::vector<int>&, int, int)':
split.cpp:130:21: warning: unused variable 'idx' [-Wunused-variable]
130 | int to = e.first, idx = e.second;
| ^~~
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |