#include "sphinx.h"
#include <cassert>
#include <numeric>
#include <set>
using namespace std;
struct DSU {
vector<int> parent;
DSU(int n) : parent(n) { iota(parent.begin(), parent.end(), 0); }
int find(int x) { return parent[x] == x ? x : parent[x] = find(parent[x]); }
void merge(int a, int b) { parent[find(a)] = find(b); }
};
vector<int> find_colours(int N, vector<int> X, vector<int> Y) {
auto only_keep = [&N] (vector<int> kept) {
vector<int> E(N, N);
for (int node : kept)
E[node] = -1;
return E;
};
vector<vector<int>> adj(N);
for (int iEdge = 0; iEdge < (int)X.size(); ++iEdge) {
int u = X[iEdge], v = Y[iEdge];
adj[u].push_back(v);
adj[v].push_back(u);
}
auto non_sphinx_comps = [&N, &adj] (vector<int> E) -> int {
vector<bool> seen(N, false);
auto Dfs = [&] (auto dfs, int node) -> void {
seen[node] = true;
for (int neighbor : adj[node])
if (E[neighbor] == N && !seen[neighbor])
dfs(dfs, neighbor);
};
int cnt_sphinx_comp = 0;
for (int node = 0; node < N; ++node) {
if (E[node] == N && !seen[node]) {
++cnt_sphinx_comp;
Dfs(Dfs, node);
}
}
return perform_experiment(E) - cnt_sphinx_comp;
};
DSU uf(N);
/// takes a subset of the neighborhood of center
/// and returns a subsubset of those neighbors st they are pairwise not
/// in the same DSU component
auto simplified = [&] (const vector<int> &orig, int center) {
set<int> dsu_roots;
vector<int> root_to_orig(N, -1);
for (int x : orig) {
root_to_orig[uf.find(x)] = x;
if (uf.find(x) != uf.find(center))
dsu_roots.insert(uf.find(x));
}
vector<int> ret;
for (int root : dsu_roots)
ret.push_back(root_to_orig[root]);
// expected: non_sphinx_comps(only_keep(ret)) == (int)ret.size())
return ret;
};
for (int center = 0; center < N; ++center) {
vector<int> less_neighbors;
for (int nei : adj[center])
if (nei < center)
less_neighbors.push_back(nei);
while (true) {
auto has_friend = [&] (vector<int> subset) { // assumes S simplified
assert(subset == simplified(subset, center));
subset.push_back(center);
return non_sphinx_comps(only_keep(subset)) < (int)subset.size();
};
auto potentials = simplified(less_neighbors, center);
if (!has_friend(potentials)) break;
while ((int)potentials.size() >= 2) {
vector<int> even, odd;
for (int i = 0; i < (int)potentials.size(); ++i) {
(i % 2 ? odd : even).push_back(potentials[i]);
}
if (has_friend(even))
potentials = even;
else
potentials = odd;
}
uf.merge(center, potentials[0]);
}
}
vector<int> answer(N, 0);
for (int node = 0; node < N; ++node) {
answer[node] = uf.find(node);
}
return answer;
}
