17B - Hierarchy - CodeForces Solution

#include <bits/stdc++.h>
using namespace std;

typedef long long ll;

struct DSU{ //Disjoint Set Union (o Union Find)
    vector <int> parent, sz;
    DSU(int n){
        parent.resize(n);
        sz.resize(n);
        for(int i=0; i<n; i++){
            parent[i] = i;
            sz[i] = 1;
        }
    }
    int find_set(int v){
        if(v == parent[v]) return v;
        return parent[v] = find_set(parent[v]);
    }
    void union_set(int a, int b){
        a = find_set(a);
        b = find_set(b);
        if(a != b){
            if(sz[a] < sz[b])
                swap(a,b);
            parent[b] = a;
            sz[a] += sz[b];
        }
    }
};

ll kruskal(int n, vector<tuple<ll,int,int>> &edges){
	DSU dsu(n);
	sort(edges.begin(), edges.end());

	ll ans = 0;
	vector<bool> rel(n,false);
	for(auto [w, u, v] : edges){ // (C++17 only)
		if(!rel[v] && dsu.find_set(u) != dsu.find_set(v)){
			ans += w;
			dsu.union_set(u, v);
			rel[v] = true;
		}
	}
	int cnt = 0;
	for(int i=0; i<n; i++) 
		if(!rel[i]) 
			cnt++;
	if(cnt > 1) 
		return -1;

	return ans;
}


int main(){
	ios_base::sync_with_stdio(0); cin.tie(0);
	int n, m;	cin >> n;
	for(int i=0; i<n; i++) cin >> m;
	cin >> m;
	vector <tuple<ll,int,int>> edges(m);
	for(int i=0; i<m; i++){
		int u, v;
		ll w;
		cin >> u >> v >> w;
		u--;
		v--;
		edges[i] = {w, u, v};
	}
	cout << kruskal(n, edges) << '\n';
	return 0;
}