Inheritance

inheritance.cpp

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