#include <bits/stdc++.h>
#include "split.h"
using namespace std;
const int N = 100'000 + 10;
int n;
pair<int, int> c[3];
vector<int> ad[N];
vector<int> ret;
int sz[N];
int low[N], num[N], it;
void initDFS(int u, int p = -1) {
sz[u] = 1;
low[u] = num[u] = ++it;
for (const auto& v : ad[u]) {
if (v == p) continue;
if (num[v]) low[u] = min(low[u], num[v]);
else {
initDFS(v, u);
sz[u] += sz[v];
low[u] = min(low[u], low[v]);
}
}
}
void doColor1(int u, int p, int idx) {
if (c[idx].first <= 0 || ret[u]) return;
c[idx].first -= 1;
ret[u] = c[idx].second;
for (const auto& v : ad[u]) {
if (v == p || num[v] < num[u]) continue;
doColor1(v, u, idx);
}
}
void doColor2(int u, int p, int idx) {
if (c[idx].first <= 0 || ret[u]) return;
c[idx].first -= 1;
ret[u] = c[idx].second;
for (const auto& v : ad[u]) {
if (v == p) continue;
doColor2(v, u, idx);
}
}
vector<int> tmp;
bool doneColor = false;
void dfs(int u, int p = - 1) {
if (doneColor) return;
for (const auto& v : ad[u]) {
if (v == p || num[v] < num[u]) continue;
if (sz[v] >= c[0].first) {
dfs(v, u);
}
if (doneColor) return;
}
tmp.clear();
int curSZ = sz[u];
for (const auto& v : ad[u]) {
if (v == p || num[v] < num[u]) continue;
if (low[v] < num[u] && curSZ - sz[v] >= c[0].first) {
curSZ -= sz[v];
} else tmp.push_back(v);
}
if (curSZ >= c[0].first && n - curSZ >= c[1].first) {
swap(ad[u], tmp);
doColor1(u, p, 0);
doColor2(p, u, 1);
doneColor = true;
} else if (curSZ >= c[1].first && n - curSZ >= c[0].first) {
swap(ad[u], tmp);
doColor1(u, p, 1);
doColor2(p, u, 0);
doneColor = true;
}
}
vector<int> find_split(int _n, int _a, int _b, int _c, vector<int> _p, vector<int> _q) {
{ // input
n = _n;
for (int i = 0; i < (int)_p.size(); ++i) {
int u = _p[i], v = _q[i];
ad[u].push_back(v);
ad[v].push_back(u);
}
c[0] = {_a, 1};
c[1] = {_b, 2};
c[2] = {_c, 3};
sort(c, c + 3);
}
ret.assign(n, 0);
initDFS(0);
dfs(0);
if (doneColor) {
for (auto& x : ret) {
if (!x) x = c[2].second, c[2].first -= 1;
}
return ret;
}
return vector<int>(n, 0);
}
| # | 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... |