#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; }