#include "split.h"
#include <bits/stdc++.h>
using namespace std;
const pair<int, int> null_pair = make_pair(-1, -1);
int n;
vector<vector<int>> t;
vector<int> abc, idx = {0, 1, 2};
vector<int> subt_n;
vector<int> split(pair<int, int> e) {
if (e == null_pair) return vector<int>(n, 0);
auto [x, y] = e;
vector<int> ret(n, 2);
int c = 0;
auto dfs = [&](const auto &self, int u, int par) -> void {
if (abc[c]) ret[u] = c, abc[c]--;
for (int v : t[u]) {
if (v == par) continue;
self(self, v, u);
}
};
dfs(dfs, x, y);
c = 1;
dfs(dfs, y, x);
for (int &x : ret) x = idx[x] + 1;
return ret;
}
vector<int> findSplit() {
subt_n.resize(n);
auto dfs1 = [&](const auto &self, int u, int par) -> void {
subt_n[u] = 1;
for (int v : t[u]) {
if (v == par) continue;
self(self, v, u);
subt_n[u] += subt_n[v];
}
};
dfs1(dfs1, 0, 0);
auto dfs2 = [&](const auto &self, int u, int par) -> pair<int, int> {
if (abc[0] <= subt_n[u] and abc[1] <= n - subt_n[u]) {
cout << "BRUH " << u << ' ' << par << ' ' << subt_n[u] << ' ' << n - subt_n[u] << endl;
return make_pair(u, par);
}
if (abc[1] <= subt_n[u] and abc[0] <= n - subt_n[u]) {
cout << "BRUH " << u << ' ' << par << ' ' << subt_n[u] << ' ' << n - subt_n[u] << endl;
return make_pair(par, u);
}
for (int v : t[u]) {
if (v == par) continue;
auto res = self(self, v, u);
if (res != null_pair) return res;
}
return null_pair;
};
return split(dfs2(dfs2, 0, 0));
}
vector<int> find_split(int N, int A, int B, int C, vector<int> U, vector<int> V) {
if (U.size() != N - 1) return vector<int>(N, 0);
n = N;
abc = vector<int>{A, B, C};
sort(idx.begin(), idx.end(), [](int i, int j) {
return abc[i] < abc[j];
});
sort(abc.begin(), abc.end());
t.resize(n);
for (int i = 0; i < n - 1; i++) {
t[U[i]].push_back(V[i]);
t[V[i]].push_back(U[i]);
}
return findSplit();
}
# | 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... |