1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59

#include <iostream>
#include <vector>
#include <numeric>

struct DSU {
    DSU(int n) {
        count_ = n;
        parent_.resize(n, 0);
        std::iota(parent_.begin(), parent_.end(), 0);
        size_.resize(n, 1);
    }
    
    int find(int a) {
        if (a != parent_[a]) {
            parent_[a] = find(parent_[a]);
        }
        return parent_[a];
    }
    
    int unite(int a, int b) {
        int root_a = find(a);
        int root_b = find(b);
        
        if (root_a != root_b) {
            if (size_[root_b] > size_[root_a]) {
                std::swap(root_b, root_a);
            }
            --count_;
            parent_[root_b] = root_a;
            size_[root_a] += size_[root_b];
        }
        
        return count_;
    }
    
private:
    std::vector<int> parent_;
    std::vector<int> size_;
    int count_;
};

class Solution {
public:
    int minCost(int n, std::vector<std::vector<int>>& edges, int k) {
        std::sort(edges.begin(), edges.end(), [](std::vector<int>& a, std::vector<int>& b) {
            return a[2] < b[2];
        });
        
        int m = edges.size();
        DSU dsu(n);
        for (int i = 0; i < m; ++i) {
            if (dsu.unite(edges[i][0], edges[i][1]) == k) {
                return edges[i][2];
            }
        }
        
        return 0;
    }
};