Solução: Inheritance
Primeiro, vamos dar a cada filho sua própria estrutura de DSU para armazenar quais cidades estão conectadas com suas ferrovias, e, em seguida, iterar pelas ferrovias em ordem decrescente de preço, dando cada ferrovia ao primeiro filho que puder aceitá-la. Isso é efetivamente o mesmo que percorrer cada filho, um de cada vez, e construir uma árvore geradora máxima para cada um deles individualmente.
No entanto, iterar linearmente por cada filho resultaria numa complexidade de 
, o que não passa nos limites. Para encontrar rapidamente o primeiro filho que pode aceitar uma ferrovia, vamos usar busca binária. Mas para ver que isso realmente funciona primeiro precisamos provar o seguinte: Independentemente da ferrovia que estivermos processando, todos os filhos após o primeiro filho que podem aceitar a ferrovia atual também podem aceitar essa ferrovia, enquanto todas as crianças antes dessa criança não podem.
Para que isso se mantenha, as componentes conectadas de cada DSU devem ser um subconjunto das posteriores. Isso é obviamente verdadeiro para quando começamos, portanto, agora tudo o que temos de provar é que dar uma ferrovia ao primeiro filho que pode aceitá-la preservará essa relação. Digamos que tenhamos dois filhos adjacentes em termos de idade: 
e 
, sendo o mais velho e o primeiro filho que pode aceitar a ferrovia atual.
Se dar a a ferrovia fizesse com que fosse impossível obter do novo , isso significaria que a ferrovia conectaria alguns dois componentes de e, portanto, também significaria que deveria ser o primeiro a aceitar a ferrovia, e não . O uso da busca binária reduz a complexidade de para 
, que passa no limite de tempo.
Veja o código a seguir para tirar eventuais dúvidas:
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| #include <bits/stdc++.h> | |
| using namespace std; | |
| const int MAXK = 2e4 + 10; | |
| const int MAXM = 3e5 + 10; | |
| struct DSU | |
| { | |
| vector<int> pai, sze; | |
| void init(int n) | |
| { | |
| pai.resize(n + 1); | |
| sze.resize(n + 1); | |
| iota(pai.begin(), pai.end(), 0); | |
| fill(sze.begin(), sze.end(), 1); | |
| } | |
| int find(int u) | |
| { | |
| return pai[u] = (pai[u] == u ? u : find(pai[u])); | |
| } | |
| void join(int u, int v) | |
| { | |
| u = find(u), v = find(v); | |
| if(u == v) return; | |
| if(sze[u] < sze[v]) swap(u, v); | |
| pai[v] = u; sze[u] += sze[v]; | |
| } | |
| bool same(int u, int v) | |
| { | |
| return (find(u) == find(v)); | |
| } | |
| } children[MAXK]; | |
| int n, m, k, resp[MAXM]; | |
| vector<tuple<int, int, int, int>> edges; // peso, id, v1, v2 | |
| int main() | |
| { | |
| ios_base::sync_with_stdio(false); | |
| cin.tie(NULL); | |
| cin >> n >> m >> k; | |
| for(int i = 1; i <= m; i++){ | |
| int u, v, w; cin >> u >> v >> w; | |
| edges.emplace_back(w, i, u, v); | |
| } | |
| // Ordeno em ordem decrescente de pesos | |
| sort(edges.begin(), edges.end(), greater<tuple<int, int, int, int>>()); | |
| for(int i = 1; i <= k; i++) | |
| children[i].init(n); | |
| for(auto [w, id, u, v] : edges){ | |
| int l = 1, r = k, opt = 0; | |
| // Busca binaria | |
| while(l <= r){ | |
| int mid = (l + r) >> 1; | |
| if(!children[mid].same(u, v)) opt = mid, r = mid - 1; | |
| else l = mid + 1; | |
| } | |
| if(opt) children[opt].join(u, v); | |
| resp[id] = opt; | |
| } | |
| for(int i = 1; i <= m; i++) | |
| cout << resp[i] << '\n'; | |
| } |