WeightedGraphCpp created by fractalbach - https://repl.it/@fractalbach/WeightedGraphCpp

main.cpp

#include <iostream>
#include <map>
#include <set>

/**
 * WeightedGraph
 * 
 * @tparam _Key:     Type of key used to reference nodes.
 *
 */
template <typename _Key>
class WeightedGraph {
public:
        WeightedGraph<_Key>() : adjmap() {};
        ~WeightedGraph<_Key>() {};
        void addNode(_Key);
        void removeNode(_Key);
        void addEdge(_Key, _Key, int);
        void removeEdge(_Key, _Key);
        void display();
        const WeightedGraph<_Key> prim();
        
protected:
        typedef _Key key_type;
        typedef std::map<_Key, std::map<_Key, int> > adjmap_type;
        typedef std::map<_Key, int> edgeSet_type;
        typedef typename adjmap_type::iterator adjmap_iter;
        typedef typename edgeSet_type::iterator edgeSet_iter;
        std::map<_Key, std::map<_Key, int> > adjmap;
};


template <typename _Key>
void WeightedGraph<_Key> :: addNode(_Key k) {
        if (adjmap.find(k) == adjmap.end()) {
                adjmap[k] = std::map<_Key, int>();
        }
}

template <typename _Key>
void WeightedGraph<_Key>::removeNode(_Key k) {
        adjmap.erase(k);
        for (adjmap_iter iter = adjmap.begin(); iter != adjmap.end(); iter++) {
                if (k == iter->first) {
                        adjmap.erase(iter);
                }
        }
}

template <typename _Key>
void WeightedGraph<_Key>::addEdge(_Key k1 , _Key k2, int w) {
        addNode(k1);
        addNode(k2);
        adjmap[k1][k2] = w;
        adjmap[k2][k1] = w;
}

template <typename _Key>
void WeightedGraph<_Key>::removeEdge(_Key k1, _Key k2) {
        if (adjmap.find(k1) != adjmap.end()) {
                adjmap[k1].erase(k2);
        }
        if (adjmap.find(k2) != adjmap.end()) {
                adjmap[k2].erase(k1);
        }
}

template <typename _Key>
void WeightedGraph<_Key>::display() {
        int totalWeight = 0;
        for (adjmap_iter it = adjmap.begin(); it != adjmap.end(); it++) {
                edgeSet_type edgeset = (it->second);
                std::cout << (it->first) << " --> {";
                if (edgeset.size() <= 0) {
                        std::cout << "}" << std::endl;
                        continue;
                }
                std::cout << "  ";
                for (edgeSet_iter it2 = edgeset.begin(); it2 != edgeset.end(); it2++) {
                        std::cout << (it2->first) << ":" << (it2->second) << ",  ";
                        totalWeight += (it2->second);
                }
                std::cout << "}" << std::endl;
        }
        std::cout << "Total Weight = " << totalWeight/2 << std::endl << std::endl;
}


/* Prim's Algorithm */

template <typename _Key>
const WeightedGraph<_Key> WeightedGraph<_Key>::prim() {
        WeightedGraph<_Key> mst;
        std::set<_Key> visited = std::set<_Key>();

        for (adjmap_iter it = adjmap.begin(); it != adjmap.end(); it++) {
                key_type current = it->first;
                edgeSet_type edgeset = it->second;
                visited.insert(current);

                if (edgeset.size() <= 0) {
                        throw "must be a connected graph to find minimum spanning tree.";
                }

                key_type minKey = edgeset.begin()->first;
                int minWeight = edgeset.begin()->second;

                for (edgeSet_iter e = edgeset.begin(); e != edgeset.end(); e++) {
                        key_type key = e->first;
                        int weight = e->second;
                        if (visited.find(key) != visited.end()) {
                                continue;
                        }
                        if (weight < minWeight) {
                                minKey = key;
                                minWeight = weight;
                        }
                }
                mst.addEdge(current, minKey, minWeight);
        }
        adjmap = mst.adjmap;
        return mst;
}




/* Compatibility */ 

template <class _Key>
class GraphWPrim : public WeightedGraph<_Key> {
public:
        GraphWPrim(_Key name) : WeightedGraph<_Key>() {}
        void addVertex(_Key k1 , _Key k2, int w) {
                WeightedGraph<_Key>::addEdge(k1, k2, w);
        }
        void print() {
                WeightedGraph<_Key>::display();
        }
        void prim() {
                WeightedGraph<_Key>::prim();
        }
};


/* Usage */

int main() {
        GraphWPrim<char> graph('A');
        graph.addVertex('B', 'A', 2);
        graph.addVertex('C', 'A', 3);
        graph.addEdge('C', 'B', 4);
        graph.addVertex('D', 'A', 3);
        graph.addEdge('C', 'D', 5);
        graph.addVertex('E', 'B', 3);
        graph.addEdge('C', 'E', 1);
        graph.addVertex('F', 'C', 6);
        graph.addEdge('D', 'F', 7);
        graph.addEdge('F', 'E', 8);
        graph.addVertex('G', 'F', 9);
        graph.print();
        graph.prim();
        graph.print();        
}