#include <bits/stdc++.h> #include "sphinx.h" using namespace std; 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]); } // 1. Vybudování BFS stromu pro získání rodičů a hloubek vector<int> parent(n, -1); vector<int> depth(n, 0); vector<int> bfs_order; vector<bool> vis(n, false); queue<int> q; q.push(0); vis[0] = true; while (!q.empty()) { int u = q.front(); q.pop(); bfs_order.push_back(u); for (int v : adj[u]) { if (!vis[v]) { vis[v] = true; parent[v] = u; depth[v] = depth[u] + 1; q.push(v); } } } vector<int> colors(n, -1); vector<int> a(n); int identified_count = 0; // 2. Hlavní smyčka přes možné barvy // Hledáme barvy pro všechny vrcholy kromě kořene (ten nemá rodiče) for (int c = 0; c < n; c++) { // Pokud jsme našli barvy pro všechny vrcholy (kromě rootu, ten řešíme na konci), končíme if (identified_count >= n - 1) break; // Rozdělíme zpracování na 3 fáze podle hloubky, abychom předešli kolizím parent-child for (int rem = 0; rem < 3; rem++) { // Helper: Zjistí počet shod v podmnožině kandidátů auto count_matches = [&](const vector<int>& candidates, int l, int r) { if (l > r) return 0; // Reset: nastavíme vše na 'n' (neutrální hodnota) fill(a.begin(), a.end(), n); int active_pairs = 0; for (int i = l; i <= r; i++) { int u = candidates[i]; int p = parent[u]; // u nastavíme na wildcard (-1) // rodiče p nastavíme na cílovou barvu c a[u] = -1; a[p] = c; // Poznámka: Jeden rodič může obsloužit více dětí najednou, to je v pořádku. // Pokud se u a p spojí, ubyde jedna komponenta. } // Pokud jsme nic nenastavili, vracíme 0 // Musíme spočítat, kolik vrcholů jsme reálně 'aktivovali' (nastavili na -1) // Protože 'active_pairs' výše jen nastavuje pole, spočteme si reálný počet unikátních aktivací int nodes_involved = 0; for (int i = l; i <= r; i++) { nodes_involved++; // Komponenta navíc vznikne jen pokud se nespojí. // Očekávaný počet komponent pro tyto páry, pokud by se nic nespojilo, je nodes_involved (každý u je sám, p je buď sám nebo s jinými). // Ale pozor: perform_experiment vrací celkový počet komponent v grafu. // Musíme porovnat s "base line". } if (nodes_involved == 0) return 0; int components = perform_experiment(a); // Každý úspěšný match (u má barvu c) sníží počet komponent o 1 oproti stavu, kdyby match nebyl. // Celkový počet vrcholů je n. // Počet komponent = n - (počet úspěšných spojení). // Toto platí, protože u je list (v našem experimentálním podgrafu, u má jen hranu na p). return n - components; }; // Smyčka pro nalezení všech výskytů barvy v aktuální vrstvě while (true) { // Vybereme kandidáty: neobarvené vrcholy v této vrstvě (kromě rootu) vector<int> candidates; for (int u : bfs_order) { if (u != 0 && colors[u] == -1 && depth[u] % 3 == rem) { candidates.push_back(u); } } if (candidates.empty()) break; // Rychlá kontrola celé skupiny int total = count_matches(candidates, 0, candidates.size() - 1); if (total == 0) break; // Žádný další vrchol v této vrstvě nemá barvu c // Binary search pro nalezení konkrétního vrcholu int l = 0, r = candidates.size() - 1; int found_idx = -1; while (l <= r) { if (l == r) { found_idx = l; break; } int mid = l + (r - l) / 2; if (count_matches(candidates, l, mid) > 0) { r = mid; } else { l = mid + 1; } } if (found_idx != -1) { int u = candidates[found_idx]; colors[u] = c; identified_count++; // Pokračujeme v cyklu while(true), abychom našli další případné vrcholy stejné barvy ve stejné vrstvě } else { break; } } } } // 3. Dořešení kořene (vrchol 0) // Kořen nemá rodiče, musíme ho porovnat s jeho sousedy, kteří už barvu mají. if (colors[0] == -1) { // Zkusíme, zda kořen nemá stejnou barvu jako některý z jeho sousedů // Iterujeme přes možné barvy c for (int c = 0; c < n; c++) { // Má kořen souseda s barvou c? bool has_neighbor_c = false; for (int v : adj[0]) { if (colors[v] == c) { has_neighbor_c = true; break; } } if (has_neighbor_c) { fill(a.begin(), a.end(), n); a[0] = -1; // Nastavíme všechny sousedy barvy c na c for (int v : adj[0]) { if (colors[v] == c) a[v] = c; } if (perform_experiment(a) < n) { colors[0] = c; break; } } } // Pokud jsme nenašli shodu s žádným sousedem, dáme mu první volnou barvu // (Většinou to bude barva, kterou cyklus výše nenašel, protože kořen byl jediný svého druhu nebo izolovaný v barvě) if (colors[0] == -1) { // Najdeme nejmenší nepoužitou barvu v okolí? Ne, prostě přiřadíme novou. // Ale musíme dodržet konzistenci, zkusíme najít první c, které jsme v grafu použili, nebo prostě volnou. // Nejbezpečnější je nechat ho být, pokud úloha vyžaduje jen partition, // ale my vracíme konkrétní čísla. Najdeme první nepoužité C? // V zadání Sphinx se obvykle vrací IDs. Pokud je root unikátní, může mít jakékoliv ID. // Pro jistotu mu dáme nějakou bezpečnou barvu (např. max_used + 1 nebo jen 0 pokud je vše 0). int max_c = -1; for(int x : colors) max_c = max(max_c, x); colors[0] = max_c + 1; } } // Fallback pro jistotu for (int i = 0; i < n; i++) { if (colors[i] == -1) colors[i] = 0; } return colors; }