test/library_checker/graph/minimum_spanning_tree.test.cpp

• View this file on GitHub
• Last update: 2024-11-18 03:54:10+09:00
• Problem: https://judge.yosupo.jp/problem/minimum_spanning_tree

Depends on

• Union-Find (lib/data_structure/union_find.hpp)

Code

#define PROBLEM "https://judge.yosupo.jp/problem/minimum_spanning_tree"
#include <bits/stdc++.h>
using namespace std;

#include "../../../lib/data_structure/union_find.hpp"

int main(){
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    int n, m; cin >> n >> m;

    using Edge = pair<int, int>;
    vector<Edge> e(m);
    vector<long long> c(m);
    for(int i = 0; i < m; i++){
        cin >> e[i].first >> e[i].second >> c[i];
    }

    vector<pair<long long, int>> sorted(m);
    for(int i = 0; i < m; i++){
        sorted[i] = {c[i], i};
    }
    UnionFind uf(n);
    sort(sorted.begin(), sorted.end());
    long long sum = 0;
    vector<int> ans;
    for(int i = 0; i < m; i++){
        auto [w, idx] = sorted[i];
        auto [u, v] = e[idx];
        if(!uf.same(u, v)){
            ans.push_back(idx);
            sum += w;
            uf.unite(u, v);
        }
    }
    cout << sum << "\n";
    for(int i = 0; i < n - 1; i++){
        if(i >= 1) cout << " ";
        cout << ans[i];
    }
    cout << "\n";
}

#line 1 "test/library_checker/graph/minimum_spanning_tree.test.cpp"
#define PROBLEM "https://judge.yosupo.jp/problem/minimum_spanning_tree"
#include <bits/stdc++.h>
using namespace std;

#line 2 "lib/data_structure/union_find.hpp"

/**
 * @brief Union-Find
 * @docs docs/data_structure/union_find.md
 */

#line 10 "lib/data_structure/union_find.hpp"

struct UnionFind{
    int V;
    std::vector<int> par;
    std::vector<int> edg;

    UnionFind(const int N) : V(N), par(N), edg(N){
        for(int i = 0; i < N; ++i){
            par[i] = -1;
            edg[i] = 0;
        }
    }

    int root(int x){
        assert(0 <= x && x < V);
        if(par[x] < 0) return x;
        return par[x] = root(par[x]);
    }

    int unite(int x, int y){
        int rx = root(x);
        int ry = root(y);
        if(rx == ry){
            edg[rx]++;
            return rx;
        }
        if(-par[rx] < -par[ry]) std::swap(rx, ry);
        par[rx] = par[rx] + par[ry];
        par[ry] = rx;
        edg[rx] += edg[ry] + 1;
        return rx;
    }

    bool same(int x, int y){
        int rx = root(x);
        int ry = root(y);
        return rx == ry;
    }

    long long size(int x){
        return -par[root(x)];
    }

    long long edge(int x){
        return edg[root(x)];
    }
};
#line 6 "test/library_checker/graph/minimum_spanning_tree.test.cpp"

int main(){
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    int n, m; cin >> n >> m;

    using Edge = pair<int, int>;
    vector<Edge> e(m);
    vector<long long> c(m);
    for(int i = 0; i < m; i++){
        cin >> e[i].first >> e[i].second >> c[i];
    }

    vector<pair<long long, int>> sorted(m);
    for(int i = 0; i < m; i++){
        sorted[i] = {c[i], i};
    }
    UnionFind uf(n);
    sort(sorted.begin(), sorted.end());
    long long sum = 0;
    vector<int> ans;
    for(int i = 0; i < m; i++){
        auto [w, idx] = sorted[i];
        auto [u, v] = e[idx];
        if(!uf.same(u, v)){
            ans.push_back(idx);
            sum += w;
            uf.unite(u, v);
        }
    }
    cout << sum << "\n";
    for(int i = 0; i < n - 1; i++){
        if(i >= 1) cout << " ";
        cout << ans[i];
    }
    cout << "\n";
}

Back to top page