library/graph/MinimumSpanningTree.hpp

View this file on GitHub
Last update: 2025-11-24 18:49:07+09:00
Include: #include "library/graph/MinimumSpanningTree.hpp"

Depends on

Verified with

test/AOJ/GRL_2_A.test.cpp

Code

#include "library/datastructure/unionfind/UnionFind.hpp"
template <typename WG, typename W = typename WG::weight_type>
std::pair<W, std::vector<int>> minimum_spanning_tree(const WG &g) {
    assert(g.is_prepared());
    int n = g.n, m = g.edges.size();
    UnionFind uf(n);
    std::vector<int> id(m);
    std::iota(id.begin(), id.end(), 0);
    std::ranges::sort(id, [&](const int i, const int j) {
        return g.edges[i].weight < g.edges[j].weight;
    });
    W res = 0;
    std::vector<int> tree;
    tree.reserve(n - 1);
    for (int i : id) {
        const auto &[from, to, weight] = g.edges[i];
        if (uf.same(from, to))
            continue;
        tree.push_back(i);
        uf.merge(from, to);
        res += weight;
    }
    assert(uf.count() == 1);
    return {res, tree};
}

#line 1 "library/datastructure/unionfind/UnionFind.hpp"
#include <numeric>
#include <vector>

class UnionFind {
    int n, num;
    std::vector<int> sz, parent;

  public:
    UnionFind() = default;
    UnionFind(int n) : n(n), num(n), sz(n, 1), parent(n, 0) {
        std::iota(parent.begin(), parent.end(), 0);
    }

    int leader(int x) {
        assert(0 <= x and x < n);
        return (x == parent[x] ? x : parent[x] = leader(parent[x]));
    }

    bool same(int x, int y) {
        assert(0 <= x and x < n and 0 <= y and y < n);
        return leader(x) == leader(y);
    }

    bool merge(int x, int y) {
        assert(0 <= x and x < n and 0 <= y and y < n);
        x = leader(x);
        y = leader(y);
        if (x == y)
            return false;
        if (sz[x] < sz[y])
            std::swap(x, y);
        sz[x] += sz[y];
        parent[y] = x;
        num--;
        return true;
    }

    int size(const int x) {
        assert(0 <= x and x < n);
        return sz[leader(x)];
    }

    int count() const { return num; }

    std::vector<std::vector<int>> groups() {
        std::vector<std::vector<int>> res(n);
        for (int i = 0; i < n; i++)
            res[leader(i)].push_back(i);
        std::erase_if(res, [](const auto &vec) { return vec.empty(); });
        return res;
    }
};
#line 2 "library/graph/MinimumSpanningTree.hpp"
template <typename WG, typename W = typename WG::weight_type>
std::pair<W, std::vector<int>> minimum_spanning_tree(const WG &g) {
    assert(g.is_prepared());
    int n = g.n, m = g.edges.size();
    UnionFind uf(n);
    std::vector<int> id(m);
    std::iota(id.begin(), id.end(), 0);
    std::ranges::sort(id, [&](const int i, const int j) {
        return g.edges[i].weight < g.edges[j].weight;
    });
    W res = 0;
    std::vector<int> tree;
    tree.reserve(n - 1);
    for (int i : id) {
        const auto &[from, to, weight] = g.edges[i];
        if (uf.same(from, to))
            continue;
        tree.push_back(i);
        uf.merge(from, to);
        res += weight;
    }
    assert(uf.count() == 1);
    return {res, tree};
}