#include "split.h"
#include <bits/stdc++.h>
using namespace std;
const int mxN = 1e5;
vector<int> t_adj[mxN];
int sz[mxN];
#define all(v) (v).begin(), (v).end()
mt19937_64 gen(chrono::high_resolution_clock::now().time_since_epoch().count());
vector<array<int, 2>> dfs_edges; // parent child
void dfs(int v, int p) {
sz[v] = 1;
for (auto u: t_adj[v]) if (u!=p) {
dfs_edges.push_back({v, u});
dfs(u, v);
sz[v] += sz[u];
}
}
void mark(int v, int p, int cnt, int clr, vector<int> &res) {
queue<array<int, 2>> q;
q.push({v, p});
while (!q.empty() && cnt) {
auto [u, p] = q.front(); q.pop();
res[u] = clr;
cnt--;
for (auto w: t_adj[u]) if (w!=p) {
q.push({w, u});
}
}
}
vector<int> try_find(int n, array<int, 3> sets, array<int, 3> ord) {
vector<int> res(n);
int x = sets[ord[0]];
int y = sets[ord[1]];
for (auto [p, c]: dfs_edges) {
int sz_c = sz[c];
int sz_p = n-sz[c];
if (sz_c>sz_p) {
swap(c, p);
swap(sz_c, sz_p);
}
if (sz_c>= x && sz_p >= y) {
mark(c, p, x, ord[0]+1, res);
mark(p, c, y, ord[1]+1, res);
for (int i=0; i<n; i++) if (!res[i]) res[i] = ord[2]+1;
break;
}
}
return res;
}
struct DSU {
vector<int> p, sz;
DSU(int n): p(n, 0), sz(n, 0) {
iota(all(p), 0);
}
int find(int v) {
if (p[v]==v) return v;
return p[v] = find(p[v]);
}
bool unite(int a, int b) {
a = find(a);
b = find(b);
if (a==b) return false;
if (sz[a]<sz[b]) swap(a, b);
p[b] = a;
sz[a] += sz[b];
return true;
}
};
void gen_spanning_tree(int n, vector<array<int, 2>> &e) {
DSU dsu(n);
shuffle(all(e), gen);
for (int i=0; i<n; i++) t_adj[i].clear();
for (auto [u, v]: e) {
if (dsu.unite(u, v)) {
t_adj[u].push_back(v);
t_adj[v].push_back(u);
}
}
dfs_edges.clear();
fill(sz, sz+n, 0);
dfs(0, -1);
}
vector<int> find_split(int n, int a, int b, int c, vector<int> p, vector<int> q) {
vector<array<int, 2>> e;
for (int i=0; i<p.size(); i++) {
int u = p[i];
int v = q[i];
e.push_back({u, v});
}
array<int, 3> sets = {a, b, c};
array<int, 3> ord = {0, 1, 2};
sort(all(ord), [&](int l, int r) {
return sets[l] < sets[r];
});
gen_spanning_tree(n, e);
return try_find(n, sets, ord);
}
# | 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... |