prim_mst.h

Go to the documentation of this file.

#include <set>

#pragma once

template<typename GraphType, typename OutputIter>
OutputIter prim_mst(const GraphType& G, OutputIter dest) {

    /* Maintain a set of edges that go out of the tree;
     * In each iteration, pick the lightest of them and include it in Tree */

    std::set<typename GraphType::edge_t> protruding_edges;
   
    GraphType Tree;

    for(auto it = G.vbegin(); it != G.vend(); it++)
        Tree.add_vertex(*it);

    const typename GraphType::vertex_type start = *G.vbegin();

    for(auto it = G.aebegin(start); it != G.aeend(start); it++)
        protruding_edges.insert(*it);

    long long sz = G.vsize();
    while(protruding_edges.size() != 0) {

        edge_t lightest = *edges.begin();
        edges.erase(*edges.begin());

        vertex_t u = lightest.either();
        vertex_t v = lightest.other(u);

        bool u_in_T = T.find(u) != T.end();
        bool v_in_T = T.find(v) != T.end();

        if(u_in_T and !v_in_T) {
            res++ = lightest;
            T.insert(v);
            for(auto it = G.aebegin(v); it != G.aeend(v); it++)
                edges.insert(*it);
        }
        if(!u_in_T and v_in_T) {
            res++ = lightest;
            T.insert(u);
            for(auto it = G.aebegin(v); it != G.aeend(v); it++)
                edges.insert(*it);
        }
    }

    assert(Tree.esize() == Tree.vsize() - 1);

    return res;
}


struct edge_t {
    long long u;
    long long v;
    long long int wt;
    edge_t(long long a, long long b, long long weight) : u(a), v(b), wt(weight) { }
    long long  either() { return u; }
    long long other(long long vertex) { return vertex == u ? v : u; }
    bool operator==(const edge_t other) const { return u == other.u and v == other.v; }
    bool operator<(const edge_t other) const { return  wt < other.wt; }
};

int main() {

    undirected_graph<long long, edge_t> g;

    long long n, m, u, v, wt;
    cin >> n >> m;
    while(m--) {
        cin >> u >> v >> wt;
        edge_t e(u, v, wt);
        g.add_edge(e);
    }

    vector<edge_t> mst;
    prim_mst(g, back_inserter(mst));

    long long cost = 0;
    for(edge_t edge : mst)
        cost += edge.wt;

    cout << cost << endl;
}

/*

concept graph {
    Iterators

    vbegin();
    vend();
    ebegin();
    eend();
    avbegin(vertex);
    avend();
    aebegin(vertex);
    aeend(vertex);

    Interface
    add_edge
    add_vertex
};

*/