A simple weighted graph in C++. #c++ #graph #weighted-graph #minimum-spanning-tree #MST

weighted-graph.hpp

#ifndef _WEIGHTED_GRAPH_HPP
#define _WEIGHTED_GRAPH_HPP 1

#include <iostream>
#include <limits>
#include <vector>
#include <set>
#include <unordered_set>
#include <queue>
#include <stdexcept>
#include <functional>
#include <cassert>

template <
    typename W = int,
    typename Enable = typename std::enable_if<std::is_arithmetic<W>::value, bool>::type>
class weighted_graph
{
public:

    typedef W weight_type;

    typedef unsigned int vertex_type;

    static const vertex_type MAX_VERTEX_NUMBER = std::numeric_limits<vertex_type>::max();

    static const vertex_type NOT_A_VERTEX = -1;

    static const weight_type DEFAULT_WEIGHT = 1u;

    static const weight_type MIN_WEIGHT = std::numeric_limits<weight_type>::min();

    static const weight_type MAX_WEIGHT = std::numeric_limits<weight_type>::max();

    struct edge
    {
        vertex_type from;
        vertex_type to;
        weight_type weight;

        edge(vertex_type u, vertex_type v, weight_type w = DEFAULT_WEIGHT):
            from(u), to(v), weight(w) {}

        bool operator<(const edge& rhs) const
        {
            if (from != rhs.from)
                return (from < rhs.from);
            else if (to != rhs.to)
                return (to < rhs.to);
            else
                return (weight < rhs.weight);
        }
    };

private:

    size_t const _n;

    std::set<edge> _e;

    void _check_vertex(vertex_type v) const
    {
        if (v >= _n)
            throw std::invalid_argument("vertex number is invalid");
    }

public:

    struct edge_iterable
    {
        friend class weighted_graph;

        typedef typename std::set<edge>::const_iterator iterator_type;

    private:

        const iterator_type _start, _end;

        edge_iterable(iterator_type start, iterator_type end):
            _start(start), _end(end) {}

    public:

        iterator_type begin() const { return _start; }

        iterator_type end() const { return _end; }
    };

    weighted_graph(size_t n): _n(n) {}

    size_t num_of_vertices() const
    { return _n; }

    size_t num_of_edges() const
    { return _e.size() / 2; }

    void add_edge(vertex_type u, vertex_type v, weight_type w)
    {
        _check_vertex(u);
        _check_vertex(v);
        _e.emplace(u, v, w);
        _e.emplace(v, u, w);
    }

    static weighted_graph read(std::istream& in)
    {
        size_t n, m;
        in >> n >> m;

        weighted_graph g(n);

        vertex_type u, v;
        weight_type w;
        for (size_t i = 0; i < m; ++i) {
            in >> u >> v >> w;
            g.add_edge(u, v, w);
        }

        return g;
    }

    /**
     * Iterate on the set of incident edges
     */
    edge_iterable adj(vertex_type from) const
    {
        auto start = _e.lower_bound(edge(from, 0u, MIN_WEIGHT));
        auto end = _e.upper_bound(edge(from, MAX_VERTEX_NUMBER, MAX_WEIGHT));
        return edge_iterable(start, end);
    };

    /**
     * Return the weight of the edge {u,v}, if it exists. Otherwise, return MAX_WEIGHT.
     */
    weight_type weight(vertex_type u, vertex_type v) const
    {
        auto it = _e.lower_bound(edge(u, v, MIN_WEIGHT));
        return (it != _e.cend() && u == it->from && v == it->to)? it->weight : MAX_WEIGHT;
    }

    //
    // Debug
    //

    void debug() const
    {
        using std::cerr;
        using std::endl;

        for (vertex_type v = 0; v < _n; ++v) {
            cerr << v << ": ";
            for (auto&& e: adj(v)) {
                cerr << e.to << "(" << e.weight << ")" << ", ";
            }
            cerr << endl;
        }
    }
};

template <typename W = int>
class minimum_spanning_tree
{
public:

    typedef W weight_type;

    typedef weighted_graph<weight_type> graph_type;

    typedef typename graph_type::vertex_type vertex_type;

    typedef typename weighted_graph<weight_type>::edge edge_type;

private:

    typedef typename std::function<bool(const edge_type&, const edge_type&)> comparator_type;

    static bool compare(const edge_type& x, const edge_type& y)
    { return x.weight > y.weight; }

    static bool reverse_compare(const edge_type& x, const edge_type& y)
    { return x.weight < y.weight; }

    const weighted_graph<weight_type>& _g;

    const comparator_type _comp;

    std::vector<vertex_type> _parent;

    weight_type _cost;

    minimum_spanning_tree(const weighted_graph<weight_type>& g, bool negate):
        _g(g), _comp(negate? reverse_compare : compare)
    {}

    void _compute()
    {
        static const vertex_type NOT_A_VERTEX = weighted_graph<weight_type>::NOT_A_VERTEX;

        const size_t N = _g.num_of_vertices();

        size_t n = 0;
        _cost = 0;
        _parent.assign(N, NOT_A_VERTEX);

        std::priority_queue<edge_type, std::vector<edge_type>, comparator_type> queue(_comp);
        for (vertex_type s = 0; s < N && n < N; ++s) {
            if (_parent[s] != NOT_A_VERTEX)
                continue;
            // Build tree with root s
            queue.emplace(s, s, 0);
            while (!queue.empty() && n < N) {
                const edge_type e = queue.top();
                const vertex_type u = e.from, v = e.to;
                assert(u == v || _parent[u] != NOT_A_VERTEX);
                queue.pop();
                if (_parent[v] != NOT_A_VERTEX)
                    continue; // already added to tree
                _parent[v] = u;
                _cost += e.weight;
                n++;
                for (const auto& p: _g.adj(v)) {
                    if (_parent[p.to] == NOT_A_VERTEX)
                        queue.push(p);
                }
            }
        }
    }

public:

    /**
     * Compute the minimum spanning tree (MST) for a graph g.
     * If @param negate is true, weights are negated (so effectively a maximum spanning tree is computed).
     */
    static minimum_spanning_tree<weight_type> of(const weighted_graph<weight_type>& g, bool negate = false)
    {
        minimum_spanning_tree<weight_type> mst(g, negate);
        mst._compute();
        return mst;
    }

    /**
     * The cost of the MST forest
     */
    weight_type cost() const
    { return _cost; }

    /**
     * The parent of each vertex in the MST forest (a root has itself as parent)
     */
    vertex_type parent(vertex_type u) const
    { return _parent.at(u); };

    /**
     * Find the lowest common ancestor (in the MST tree) of two vertices u and v.
     * If u and v are not connected (belong to different trees of the MST forest), return NOT_A_VERTEX.
     */
    vertex_type lowest_common_ancestor(vertex_type u, vertex_type v) const
    {
        static const vertex_type NOT_A_VERTEX = weighted_graph<weight_type>::NOT_A_VERTEX;

        const size_t N = _g.num_of_vertices();
        if (u >= N || v >= N)
            throw out_of_range("invalid vertex number");

        if (u == v)
            return u;

        std::unordered_set<vertex_type> a; // ancestors of u

        vertex_type u1 = u;
        do {
            u = u1;
            a.insert(u);
            u1 = _parent[u];
        } while (u != u1);

        vertex_type v1 = v;
        do {
            v = v1;
            if (a.count(v) > 0) {
                return v;
            }
            v1 = _parent[v];
        } while (v != v1);

        return NOT_A_VERTEX;
    }
};

#endif // _WEIGHTED_GRAPH_HPP