Inheritance

Explicación

A primera vista, parece que Kruskal sería más apropiado para este problema, así que intentémoslo.

Primero demos a cada hijo su propia estructura de datos DSU para guardar qué ciudades están conectadas con sus ferrocarriles, y luego iteremos por los ferrocarriles en orden decreciente de precio, dando cada ferrocarril al primer hijo que pueda aceptarlo. Esto es efectivamente lo mismo que recorrer a cada hijo de a uno y construir un árbol de expansión máxima para cada uno de ellos individualmente.

Sin embargo, iterar linealmente por toda la lista de DSU resultaría en un tiempo de ejecución de alrededor de $\mathcal{O}(MK \alpha(N))$, que es demasiado lento dadas las cotas.

Para hallar rápido el primer hijo que puede aceptar un ferrocarril, ¡intentemos usar búsqueda binaria! Pero para que la búsqueda binaria realmente funcione, primero necesitamos demostrar lo siguiente:

No importa qué ferrocarril estemos procesando, todos los hijos posteriores al primer hijo que puede aceptar el ferrocarril actual también pueden aceptar dicho ferrocarril, mientras que todos los hijos anteriores a ese hijo no pueden.

Para que esto se cumpla, las componentes conexas de cada DSU deben poder obtenerse a partir de las del siguiente agregando algunas aristas (cero aristas también está bien). Esto es obviamente cierto al empezar, así que ahora solo tenemos que demostrar que dar un ferrocarril al primer hijo que puede aceptarlo preservará esta relación.

Digamos que teníamos dos hijos adyacentes en antigüedad: $c_i$ y $c_{i+1}$, con $c_i$ siendo el mayor, y $c_{i+1}$ siendo el primer hijo que puede aceptar el ferrocarril actual. Si dar a $c_{i+1}$ el ferrocarril haría imposible obtener $c_i$ a partir del nuevo $c_{i+1}$, entonces eso significaría que el ferrocarril conectaría alguna pareja de componentes de $c_i$ y por lo tanto también significaría que $c_i$ debería ser el primero en aceptar el ferrocarril, no $c_{i+1}$.

Usar búsqueda binaria reduce $\mathcal{O}(MK \alpha(N))$ a $\mathcal{O}(M \log K \alpha(N))$, que corre a tiempo.

Implementación

Complejidad temporal: $\mathcal{O}(NK + M \log K \alpha(N))$

#include <algorithm>
#include <iostream>
#include <vector>
 
using std::cout;
using std::endl;
using std::vector;
 
// BeginCodeSnip{DSU}
class DisjointSets {
  private:
	vector<int> parents;
	vector<int> sizes;
 
  public:
	DisjointSets(int size) : parents(size), sizes(size, 1) {
		for (int i = 0; i < size; i++) { parents[i] = i; }
	}
 
	int get_ultimate(int n) {
		return parents[n] == n ? n : (parents[n] = get_ultimate(parents[n]));
	}
 
	int same_set(int n1, int n2) { return get_ultimate(n1) == get_ultimate(n2); }
 
	bool link(int n1, int n2) {
		n1 = get_ultimate(n1);
		n2 = get_ultimate(n2);
		if (n1 == n2) { return false; }
		if (sizes[n1] < sizes[n2]) { std::swap(n1, n2); }
		sizes[n1] += sizes[n2];
		parents[n2] = n1;
		return true;
	}
};
// EndCodeSnip
 
struct Railroad {
	int id;
	int city1;
	int city2;
	int cost;
};
 
int main() {
	int city_num;
	int rail_num;
	int child_num;
	std::cin >> city_num >> rail_num >> child_num;
 
	vector<Railroad> railroads(rail_num);
	for (int r = 0; r < rail_num; r++) {
		Railroad &rr = railroads[r];
		rr.id = r;
		std::cin >> rr.city1 >> rr.city2 >> rr.cost;
		rr.city1--;
		rr.city2--;
	}
	std::sort(railroads.begin(), railroads.end(),
	          [](const Railroad &r1, const Railroad &r2) { return r1.cost > r2.cost; });
 
	vector<DisjointSets> children(child_num, DisjointSets(city_num));
	vector<int> ownership(rail_num);
	for (int r = 0; r < rail_num; r++) {
		const Railroad &rr = railroads[r];
		int lo = 0;
		int hi = child_num - 1;
		int valid = -1;
		while (lo <= hi) {
			int mid = (lo + hi) / 2;
			if (!children[mid].same_set(rr.city1, rr.city2)) {
				valid = mid;
				hi = mid - 1;
			} else {
				lo = mid + 1;
			}
		}
 
		if (valid != -1) { children[valid].link(rr.city1, rr.city2); }
		ownership[rr.id] = valid + 1;
	}
	for (int c : ownership) { cout << c << '\n'; }
}