#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); }
};
class vgraph {
public:
vgraph(const vector<vector<int>> &orig_adj, DSU &uf) {
orig_N = orig_adj.size();
vector<int> root_to_vnode(orig_N, -1);
for (int o_node = 0; o_node < orig_N; ++o_node) {
if (o_node == uf.find(o_node))
root_to_vnode[o_node] = vN++;
}
vadj.resize(vN);
bags.resize(vN);
auto virtualize = [&] (int node) {
return root_to_vnode[uf.find(node)];
};
for (int o1 = 0; o1 < orig_N; ++o1) {
int v1 = virtualize(o1);
bags[v1].push_back(o1);
for (int o2 : orig_adj[o1]) {
int v2 = virtualize(o2);
if (v1 == v2) continue;
vadj[v1].push_back(v2);
vadj[v2].push_back(v1);
}
}
for (auto &row : vadj) {
sort(row.begin(), row.end());
row.resize(unique(row.begin(), row.end()) - row.begin());
}
}
int vN = 0;
vector<vector<int>> vadj;
vector<int> expand(const vector<int> &v) {
vector<int> ret(orig_N);
for (int virt = 0; virt < vN; ++virt) {
for (int orig : bags[virt])
ret[orig] = v[virt];
}
return ret;
}
int perform(vector<int> vE) {
return perform_experiment(expand(vE));
}
private:
int orig_N;
vector<vector<int>> bags;
};
int cnt_known_comps(const vector<vector<int>> &adj, const vector<int> &E) {
int N = adj.size();
int cnt_known_comps = 0;
vector<bool> seen(N, false);
auto dfs = [&] (auto dfs, int node) -> void {
seen[node] = true;
for (int neighbor : adj[node])
if (E[node] == E[neighbor] && !seen[neighbor])
dfs(dfs, neighbor);
};
for (int node = 0; node < N; ++node) {
if (E[node] != -1 && !seen[node]) {
++cnt_known_comps;
dfs(dfs, node);
}
}
return cnt_known_comps;
}
/// in the given graph, each edge must have endpoints with distinct colors
vector<int> phase_two(vgraph &graph, int nb_colors) {
int vN = graph.vN;
auto &vadj = graph.vadj;
/// every node in search must have a neighbor in helpers
/// returns true if there is a node in search with given color
auto exists_node = [&] (int color, const vector<int> &search, const vector<int> &helpers) {
vector<int> vE(vN, nb_colors);
for (int node : search)
vE[node] = -1;
for (int node : helpers)
vE[node] = color;
int cnt_unknown = graph.perform(vE) - cnt_known_comps(vadj, vE);
// at least one search attached to a helper?
return cnt_unknown < (int)search.size();
};
/// assumes exists_node is true, returns a node in search with color
auto find_one_node = [&] (int color, vector<int> search, const vector<int> &helpers) {
while ((int)search.size() >= 2) {
vector<int> s1, s2;
for (int node : search) {
s1.push_back(node);
swap(s1, s2);
}
auto chosen = (exists_node(color, s1, helpers)) ? s1 : s2;
search = chosen;
}
return search[0];
};
vector<int> answers(vN, -1);
auto mark_all_nodes = [&] (int color, vector<int> search, const vector<int> &helpers) {
while (exists_node(color, search, helpers)) {
int node = find_one_node(color, search, helpers);
answers[node] = color;
search.erase(find(search.begin(), search.end(), node));
}
};
vector<int> even, odd;
{
vector<bool> seen(vN, false);
auto dfs = [&] (auto dfs, int node, int depth) -> void {
seen[node] = true;
(depth % 2 ? odd : even).push_back(node);
for (int neighbor : vadj[node])
if (!seen[neighbor])
dfs(dfs, neighbor, depth+1);
};
dfs(dfs, 0, 0);
}
for (int color = 0; color < nb_colors; ++color) {
mark_all_nodes(color, even, odd);
mark_all_nodes(color, odd, even);
}
return graph.expand(answers);
}
vector<int> find_colours(int N, vector<int> X, vector<int> Y) {
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);
}
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
vector<int> E(N, N);
for (int node : subset)
E[node] = -1;
E[center] = -1;
int cnt_unknown = perform_experiment(E) - cnt_known_comps(adj, E);
return cnt_unknown < (int)subset.size() + 1;
};
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]);
}
}
vgraph compressed(adj, uf);
if (compressed.vN == 1) {
int only_color = 0;
while (true) {
vector<int> E(N, -1);
E[0] = only_color;
if (perform_experiment(E) == 1) break;
++only_color;
}
return vector<int>(N, only_color);
}
return phase_two(compressed, N);
}
