/**
* [S1-3] First 3 subtasks are pretty straight-forward. [S1] is either a cycle or
* a chain, which is simple. For [S2] we can just find a spanning tree and pick
* a random leaf node for set A. [S3] is just a DFS on tree, which means find
* a spanning tree + DFS scores 40. Don't have much thoughts on [S4-5].
*
* Time Complexity: O(n + m * log(m))
* Implementation 1 (only solves S1-3)
*/
#include <bits/stdc++.h>
#include "split.h"
typedef std::vector<int> vec;
vec parent, rank;
inline int find(int k) {
if (parent[k] == k)
return k;
return parent[k] = find(parent[k]);
}
bool merge(int a, int b) {
a = find(a), b = find(b);
if (a == b)
return false;
if (rank[b] > rank[a])
std::swap(a, b);
parent[b] = a, rank[a] += rank[b];
return true;
}
struct set_t {
int size, original_idx;
};
std::vector<vec> tree;
vec subtree;
void find_subsize(int k, int parent) {
subtree[k] = 1;
for (int child : tree[k]) {
if (child == parent)
continue;
find_subsize(child, k);
subtree[k] += subtree[child];
}
}
vec find_split(int n, int a, int b, int c, vec p, vec q) {
int m = p.size();
std::vector<set_t> sizes(3);
sizes[0].size = a, sizes[1].size = b, sizes[2].size = c;
for (int i = 0; i < 3; i++)
sizes[i].original_idx = i + 1;
std::sort(sizes.begin(), sizes.end(),
[](const set_t& s1, const set_t& s2) {
return s1.size < s2.size;
});
parent.resize(n);
rank.resize(n);
for (int i = 0; i < n; i++)
parent[i] = i, rank[i] = 1;
std::vector<bool> tree_edge(m, false);
tree.assign(n, vec());
subtree.resize(n);
for (int k = 0; k < m; k++) {
if (merge(p[k], q[k])) {
tree[p[k]].emplace_back(q[k]);
tree[q[k]].emplace_back(p[k]);
} else {
tree_edge[k] = true;
}
}
find_subsize(0, -1);
int node1 = -1, node2 = -1;
bool found = false;
for (int k = 0; k < m && !found; k++) {
if (tree_edge[k])
continue;
node1 = p[k], node2 = q[k];
if (subtree[node1] > subtree[node2])
std::swap(node1, node2);
int size1 = subtree[node1], size2 = n - size1;
if (size1 > size2) {
std::swap(size1, size2);
std::swap(node1, node2);
}
if (size1 >= sizes[0].size && size2 >= sizes[1].size)
found = true;
}
if (!found)
return vec(n, 0);
std::cerr << node1 << ": " << sizes[0].size << ", " << node2 << ": " << sizes[1].size << std::endl;
vec ans(n, 2);
std::vector<bool> visited(n, false);
visited[node1] = visited[node2] = true;
std::queue<int> bfs_queue;
bfs_queue.emplace(node1);
for (int _ = 0; _ < sizes[0].size; _++) {
assert(!bfs_queue.empty());
int t = bfs_queue.front();
bfs_queue.pop();
ans[t] = 0;
for (int neighb : tree[t]) {
if (visited[neighb])
continue;
visited[neighb] = true;
bfs_queue.emplace(neighb);
}
}
bfs_queue = std::queue<int>();
bfs_queue.emplace(node2);
for (int _ = 0; _ < sizes[1].size; _++) {
assert(!bfs_queue.empty());
int t = bfs_queue.front();
bfs_queue.pop();
ans[t] = 1;
for (int neighb : tree[t]) {
if (visited[neighb])
continue;
visited[neighb] = true;
bfs_queue.emplace(neighb);
}
}
for (int i = 0; i < n; i++)
ans[i] = sizes[ans[i]].original_idx; // convert back to (1,2,3) - (A,B,C)
return ans;
}
