Revise paste Untitled

#include <iostream>
using namespace std;
const int maxn = 1e5 + 10;
int idx[maxn];
int sz[maxn];
void init() {
    for(int i = 0; i < maxn; i++) {
        idx[i] = i;
        sz[i] = 1;
    }
}

int find_root(int A) {
    while(idx[A] != A) {
        idx[A] = idx[idx[A]];
        A = idx[A];
    }
    return A;
}
void unite(int A, int B) {
    int root_A = find_root(A);
    int root_B = find_root(B);
    if(root_A != root_B) {
        if(sz[root_A] < sz[root_B]) {
            sz[root_B] += sz[root_A];
            idx[root_A] = idx[root_B];
        }
        else {
            idx[root_B] = idx[root_A];
            sz[root_A] += sz[root_B];
        }
    }
}
bool check_if_they_belong_to_same_set(int A, int B) {
    return find_root(A) == find_root(B);
}
int main()
{
    init();
    int n, m;
    cin >> n >> m;
    
    vector<pair<int, pair<int, int>>> graph;
    
    for(int i = 0; i < m; i++) {
        int a, b, c;
        cin >> a >> b >> c;
        
        graph.push_back(make_pair(c, make_pair(a, b)));
    }
    sort(graph.begin(), graph.end());
    int res = 0;
    for(int i = 0; i < m; i++) {
        int a = graph[i].second.first;
        int b = graph[i].second.second;
        int c = graph[i].first;
        if(!check_if_they_belong_to_same_set(a, b)) {
            unite(a, b);
            res += c;
        }
    }
    cout << res << endl;
    
}