#include "split.h"
#include <bits/stdc++.h>
using namespace std;
class DisjointSet {
private:
vector<int> p, sz;
public:
DisjointSet(int n) {
p.assign(n, 0);
iota(begin(p), end(p), 0);
sz.assign(n, 1);
}
int Find(int x) {
return p[x] == x ? x : p[x] = Find(p[x]);
}
bool Unite(int x, int y) {
x = Find(x), y = Find(y);
if (x == y) return false;
if (sz[x] > sz[y]) swap(x, y);
p[x] = y;
sz[y] += sz[x];
}
int Size(int x) {
return sz[Find(x)];
}
bool SameSet(int x, int y) {
return Find(x) == Find(y);
}
};
int find_root(int mxsize, const vector<vector<int>> &adj) {
vector<int> sz(adj.size());
vector<int> par(adj.size(), -1);
function<void(int, int)> dfs = [&](int u, int p) {
sz[u] = 1;
for (auto v : adj[u]) if (v != par[u]) {
par[v] = u;
dfs(v, u);
sz[u] += sz[v];
}
};
dfs(0, -1);
int u = 0;
while (true) {
pair<int, int> mx = {-1, -1};
for (auto v : adj[u]) if (v != par[u]) {
mx = max(mx, {sz[v], v});
}
if (mx.first <= mxsize) return u;
u = mx.second;
}
}
vector<int> find_split(int n, int a, int b, int c,
vector<int> p, vector<int> q) {
DisjointSet d(n);
vector<int> splits(n, 0);
int m = p.size();
int ca, cb, cc; // color of a, b, c
{ // recolor a, b, c such that a <= b <= c holds
vector<pair<int, int>> split_sizes = {{a, 1}, {b, 2}, {c, 3}};
sort(begin(split_sizes), end(split_sizes));
tie(a, ca) = split_sizes[0];
tie(b, cb) = split_sizes[1];
tie(c, cc) = split_sizes[2];
}
vector<vector<int>> adj(n); // spanning tree adjacency list
for (int i = 0; i < m; i++) {
if (d.Unite(p[i], q[i])) {
adj[p[i]].emplace_back(q[i]);
adj[q[i]].emplace_back(p[i]);
}
}
d = DisjointSet(n);
int r = find_root(n - a, adj);
// Unite all subtrees of children of root
for (int i = 0; i < n; i++) {
for (auto j : adj[i]) {
if (i == r || j == r) continue;
d.Unite(i, j);
}
}
function<void(int, int&, int, int)> color = [&](int u, int &n, int c, int bad) {
if (u == bad || splits[u] != 0 || n == 0) return;
splits[u] = c, n--;
for (auto v : adj[u]) color(v, n, c, bad);
};
for (int i = 0; i < m; i++) {
if (d.SameSet(p[i], q[i])) continue;
if (p[i] == r || q[i] == r) continue;
if (d.Size(p[i]) > d.Size(q[i])) swap(p[i], q[i]);
if (d.Size(q[i]) >= a) { // The subtree where q[i] is in is for set A
for (auto v : adj[r]) if (d.SameSet(v, q[i])) {
color(v, a, ca, r); // color merged subtree for set A
break;
}
color(r, b, cb, -1); // color root and its surrounding subgraph for set B
break;
} else {
d.Unite(p[i], q[i]);
adj[p[i]].emplace_back(q[i]);
adj[q[i]].emplace_back(p[i]);
}
}
if (splits[r] != 0) { // if there exists a solution, we put nodes not in A or B to set C
replace(begin(splits), end(splits), 0, cc);
}
return splits;
}
Compilation message
split.cpp: In member function 'bool DisjointSet::Unite(int, int)':
split.cpp:26:3: warning: control reaches end of non-void function [-Wreturn-type]
}
^
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Incorrect |
5 ms |
256 KB |
jury found a solution, contestant did not |
| 2 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Incorrect |
5 ms |
384 KB |
invalid split: #1=0, #2=1, #3=2 |
| 2 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Incorrect |
5 ms |
256 KB |
invalid split: #1=1, #2=0, #3=4 |
| 2 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Incorrect |
5 ms |
384 KB |
invalid split: #1=8, #2=0, #3=1 |
| 2 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Incorrect |
5 ms |
256 KB |
jury found a solution, contestant did not |
| 2 |
Halted |
0 ms |
0 KB |
- |
