Graph/mst.h

• View this file on GitHub
• Last update: 2024-03-17 02:51:32+00:00

Depends on

• DataStructure/DSU/DisjointSet.h

Verified with

• Graph/tests/aizu_grl_2_a_mst.test.cpp
• Graph/tests/yosupo_mst.test.cpp

Code

// MST. 0-based index
//
// Returns:
// {mst cost, edges in mst}
//
// If graph is not connected, returns forest (number of edges will be < n-1)

#include "../DataStructure/DSU/DisjointSet.h"

// MST {{{
using ll = long long;
template<typename EdgeT>
std::pair<ll, std::vector<EdgeT>> mst(
        int n,
        std::vector<EdgeT> edges) {
    std::sort(edges.begin(), edges.end());

    DSU dsu(n + 1);  // tolerate 1-based index
    ll total = 0;
    vector<EdgeT> tree;
    for (const auto& e : edges) {
        if (dsu.merge(e.u, e.v)) {
            total += e.c;
            tree.push_back(e);
        }
    }
    return {total, tree};
}
struct Edge {
    int u, v;
    ll c;
};
bool operator < (const Edge& a, const Edge& b) {
    return a.c < b.c;
}
ostream& operator << (ostream& out, const Edge& e) {
    out << e.u << " - " << e.v << " [" << e.c << ']';
    return out;
}
// }}}

#line 1 "Graph/mst.h"
// MST. 0-based index
//
// Returns:
// {mst cost, edges in mst}
//
// If graph is not connected, returns forest (number of edges will be < n-1)

#line 1 "DataStructure/DSU/DisjointSet.h"
// DisjointSet {{{
struct DSU {
    vector<int> lab;

    DSU(int n) : lab(n+1, -1) {}

    int getRoot(int u) {
        if (lab[u] < 0) return u;
        return lab[u] = getRoot(lab[u]);
    }

    bool merge(int u, int v) {
        u = getRoot(u); v = getRoot(v);
        if (u == v) return false;
        if (lab[u] > lab[v]) swap(u, v);
        lab[u] += lab[v];
        lab[v] = u;
        return true;
    }

    bool same_component(int u, int v) {
        return getRoot(u) == getRoot(v);
    }

    int component_size(int u) {
        return -lab[getRoot(u)];
    }
};
// }}}
#line 9 "Graph/mst.h"

// MST {{{
using ll = long long;
template<typename EdgeT>
std::pair<ll, std::vector<EdgeT>> mst(
        int n,
        std::vector<EdgeT> edges) {
    std::sort(edges.begin(), edges.end());

    DSU dsu(n + 1);  // tolerate 1-based index
    ll total = 0;
    vector<EdgeT> tree;
    for (const auto& e : edges) {
        if (dsu.merge(e.u, e.v)) {
            total += e.c;
            tree.push_back(e);
        }
    }
    return {total, tree};
}
struct Edge {
    int u, v;
    ll c;
};
bool operator < (const Edge& a, const Edge& b) {
    return a.c < b.c;
}
ostream& operator << (ostream& out, const Edge& e) {
    out << e.u << " - " << e.v << " [" << e.c << ']';
    return out;
}
// }}}

Back to top page