#include <bits/stdc++.h>
#include "sphinx.h"
using namespace std;
// Globální proměnné pro DSU (Disjoint Set Union) pro správu komponent
vector<int> parent;
int find_set(int v) {
if (v == parent[v]) return v;
return parent[v] = find_set(parent[v]);
}
void union_sets(int a, int b) {
a = find_set(a);
b = find_set(b);
if (a != b) parent[b] = a;
}
std::vector<int> find_colours(int N, std::vector<int> X, std::vector<int> Y) {
int n = N;
vector<vector<int>> adj(n);
for (size_t i = 0; i < X.size(); i++) {
adj[X[i]].push_back(Y[i]);
adj[Y[i]].push_back(X[i]);
}
// colors[i] bude držet ID barvy. Na začátku -1 (neznámá).
vector<int> colors(n, -1);
// Inicializace DSU
parent.resize(n);
iota(parent.begin(), parent.end(), 0);
// Pole pro dotazy
vector<int> a(n);
for (int i = 0; i < n; i++) {
// 1. Zjistíme, jaké barvy mají sousedé vrcholu i, kteří už jsou zpracovaní
vector<int> neighbor_colors;
for (int neighbor : adj[i]) {
if (neighbor < i) { // Sousedé, které už jsme prošli
neighbor_colors.push_back(find_set(neighbor));
}
}
// Odstraníme duplicity a seřadíme, abychom mohli binárně vyhledávat
sort(neighbor_colors.begin(), neighbor_colors.end());
neighbor_colors.erase(unique(neighbor_colors.begin(), neighbor_colors.end()), neighbor_colors.end());
int found_color = -1; // -1 znamená "nová barva"
// 2. Binární vyhledávání pouze nad barvami sousedů
// Hledáme, jestli i patří do některé z existujících barevných komponent sousedů
int l = 0, r = (int)neighbor_colors.size() - 1;
while (l <= r) {
int mid = l + (r - l) / 2;
// Nastavíme dotaz:
// Vrchol i bude mít barvu -1
// Sousedé s barvami v rozsahu [l, mid] budou mít taky -1
// Všichni ostatní budou mít svou unikátní barvu (nebo N, aby se nepletli)
fill(a.begin(), a.end(), n); // Reset na "ignorovat"
a[i] = -1;
// Vybereme sadu barev, kterou testujeme
vector<int> target_colors;
for(int k = l; k <= mid; k++) target_colors.push_back(neighbor_colors[k]);
// Označíme všechny vrcholy, které mají tuto barvu
// (poznámka: pro optimalizaci bychom si mohli držet seznamy vrcholů pro každou barvu)
int nodes_in_subset = 0;
for(int u = 0; u < i; u++) {
int root = find_set(u);
bool matches = false;
// Je u součástí testované množiny barev?
// Protože target_colors je malá a setříděná, můžeme použít binary_search nebo prostý průchod
for(int c : target_colors) if(c == root) matches = true;
if(matches) {
a[u] = -1;
nodes_in_subset++;
} else {
a[u] = u; // Každý jiný vrchol má svou unikátní barvu, aby se nespojil
}
}
// Teorie: Pokud se 'i' spojí s testovanou množinou, počet komponent bude o 1 menší,
// než kdyby byl 'i' izolovaný.
// Počet komponent tvořených vrcholy v podmnožině + vrchol i (všechny mají -1)
// Pokud se i NEPOJÍ, bude to 1 (celá podmnožina se slije) + 1 (vrchol i) = 2 komponenty z pohledu -1 barvy?
// Ne, jednodušší je porovnat s perform_experiment.
// Mnohem robustnější logika pro Sphinx:
// Pokud a[i] = -1 a skupina S má a[v] = -1.
// Pokud se i spojí s S, celkový počet komponent v grafu klesne o 1 oproti stavu, kdy by i mělo unikátní barvu.
int res = perform_experiment(a);
// Kolik komponent bychom čekali?
// Všichni v 'a' s hodnotou -1 by se měli slít do jedné komponenty, POKUD jsou propojení přes staré hrany.
// Ale my víme, že staré hrany už barvy spojily. Takže všechny vrcholy s barvami v target_colors
// už tvoří souvislý celek (nebo více) v rámci -1.
// Trik: prostě se podíváme, jestli se 'i' chytilo.
// Zjednodušení pro tuto úlohu:
// Počet komponent = (počet vrcholů s unikátní barvou) + (počet komponent vzniklých z -1).
// Vrcholy s a[u]=u přispívají každý 1.
// Vrcholy s -1: Pokud se 'i' spojí s někým, bude to 1 komponenta. Pokud ne, budou to 1 (skupina) + 1 (vrchol i) = 2?
// Pozor, skupina target_colors nemusí být souvislá mezi sebou navzájem,
// ale my testujeme jestli se 'i' připojí k ALESPOŇ JEDNOMU z nich.
// Uděláme to jednodušeji podle definice Sphinx:
// Očekávaný počet = (počet vrcholů 0..i-1 mimo target_colors) + (počet disjunktních komponent v target_colors) + 1 (vrchol i).
// Pokud real < očekávaný, znamená to, že se i spojilo s někým.
int expected = 0;
// Spočítáme očekávané komponenty ručně nebo s DSU (rychlejší):
// Protože dáváme a[u]=u pro ostatní, každý je komponenta.
int distinct_others = i - nodes_in_subset;
// Komponenty uvnitř target_colors:
// Protože target_colors jsou už ID z DSU, každá barva je jedna komponenta.
// A všechny jsme přebarvili na -1. Orákulum je spojí jen pokud mezi nimi vedou hrany (to už by byly spojené)
// nebo přes vrchol i.
// Takže očekáváme (počet barev v target_colors) + 1 (vrchol i).
// Pokud se i připojí k některé barvě, klesne to.
int components_in_target = target_colors.size();
int total_expected = distinct_others + components_in_target + 1;
if (res < total_expected) {
// Spojilo se to! Hledaná barva je v levé polovině.
found_color = target_colors.back(); // Prozatím uložíme
// Musíme pokračovat v hledání, abychom našli tu první/konkrétní
if (l == mid) {
found_color = target_colors[0]; // Našli jsme přesně
break;
}
r = mid;
} else {
// Nespojilo se, barva je v pravé polovině
l = mid + 1;
}
}
if (found_color != -1) {
// i má stejnou barvu jako found_color
union_sets(found_color, i);
colors[i] = find_set(i); // Uložíme reprezentanta
} else {
// i má novou unikátní barvu
// V DSU už je i samo sebou, což je ok.
colors[i] = i;
}
}
// Normalizace barev na 0..K (volitelné, Sphinx obvykle bere jakákoliv čísla)
return colors;
}
