Implementation of Prim's algorithm in C++

Question

Could anyone comment on what could be done better and if I made any mistakes?

#include <iostream>
#include <vector>
#include <algorithm>


using namespace std;

typedef pair<pair<int, int>, int> Edge; //a, b, length
// A---------B
//   length


vector<Edge> prims_algo(vector<vector<int>> graph, int number_of_nodes){
    vector<int> unvisited;
    vector<int> visited;
    vector<Edge> result;

    //mark first as visited and mark the rest as unvisited
    for (int i = 1; i < number_of_nodes; i++)
        unvisited.push_back(i);
    visited.push_back(0);

    while (!unvisited.empty()) {

        vector<Edge> edges_with_lengths;
        //put all edges (with their lengths) from nodes that are in visited
        for(auto node : visited) {
            for (int sec_node = 0; sec_node < number_of_nodes; sec_node++) {
                if (graph[node][sec_node] > 0 && find(unvisited.begin(), unvisited.end(), sec_node) != unvisited.end()) {
                    //add if there is connection and second node is not visited yet
                    Edge e = make_pair(make_pair(node, sec_node), graph[node][sec_node]);
                    edges_with_lengths.push_back(e);
                }
            }
        }

        //find the shortest edge
        pair<pair<int, int>, int> the_shortest;
        the_shortest = edges_with_lengths.front();
        for(auto i: edges_with_lengths){
            if(the_shortest.second > i.second)
                the_shortest = i;
        }

        //add the shortest path to the result
        result.push_back(the_shortest);

        //remove a node that the shortest edge leads to
        unvisited.erase(remove(unvisited.begin(), unvisited.end(), the_shortest.first.second), unvisited.end());

        //mark this node as visited
        visited.push_back(the_shortest.first.second);
    };
    return result;
}


int main() {
    /*
       2    3
   (0)--(1)--(2)
    |   / \   |
   6| 8/   \5 |7
    | /     \ |
   (3)-------(4)
         9          */
    vector<vector<int>> graph =
                     {{{0, 2, 0, 6, 0},
                       {2, 0, 3, 8, 5},
                       {0, 3, 0, 0, 7},
                       {6, 8, 0, 0, 9},
                       {0, 5, 7, 9, 0},
    }};

    vector<Edge> result = prims_algo(graph, 5);
    for(auto i:result){
        cout<<"EDGE [ " <<i.first.first<<", "<<i.first.second<<"], length: "<<i.second<<endl;
    }

    return 0;
}

Answer 1

Code looks OK.

I would suggest you to do some class instead of Edge:

//typedef pair<pair<int, int>, int> Edge; //a, b, length
struct Edge{
   int a;
   int b;
   int length;
};

Also I suggest to have some predefined type for vector<Edge>

using EdgeContainer = vector<Edge>;

You have duplication near //find the shortest edge, you probably need Edge there

Here the_shortest is probably not initialized. Also you probably might use something from <algorithm> (but I will do it in same way as you).

    //find the shortest edge
    pair<pair<int, int>, int> the_shortest;
    the_shortest = edges_with_lengths.front();
    for(auto i: edges_with_lengths){
        if(the_shortest.second > i.second)
            the_shortest = i;
    }

Another probably micro optimization would be to write ++i in loops:

//for (int i = 1; i < number_of_nodes; i++)
for (int i = 1; i < number_of_nodes; ++i){
   // ...
}

Also in C++11 main() does not need to return 0;.

Answer 2

Consider naming the function for the result it produces rather than the algorithm used. The user of the function generally doesn't care how the result is found, and will not mind if we change to a different algorithm in future versions, as long as it returns the minimum spanning tree of the input nodes.

Avoid using namespace std; - that's a very large (and changing) namespace, and it's much safer and clearer to use qualified names where you need them.

We are missing an include of <utility> for the use of std::pair. That said, the nested pair looks like a candidate for std::tuple or, more likely, a struct with proper names for its elements.

The function doesn't need to modify the contents of graph, so there's no need to make a copy; we should pass by const-ref instead. Do we really need the number_of_nodes argument? I think that's just the inner vector's length.

We could use std::iota instead of a loop to populate unvisited:

#include <numeric>

    unvisited.resize(number_of_nodes - 1);
    std::iota(unvisited.begin(), unvisited.end(), 1);

Similarly, for finding the shortest edge, we could use std::min_element() instead of our own loop. Or we could keep track of the minimum as we examine the lengths, instead of populating a vector and then searching it.

I think that a std::set would be a more suitable storage for unvisited. That would enable more efficient membership test and removal. The result and visited vectors can be given sufficient capacity for the number of elements they will hold at the end using reserve.

---

Here's what I get after making these improvements:

#include <algorithm>
#include <climits>
#include <set>
#include <vector>

struct Edge
{
    // A---------B
    //   length
    int a;
    int b;
    int length;
};

std::vector<Edge> minimum_spanning_tree(const std::vector<std::vector<int>>& graph)
{
    std::set<int> unvisited;
    std::vector<int> visited;
    std::vector<Edge> result;

    auto const number_of_nodes  = graph.size();

    // mark first as visited and mark the rest as unvisited
    visited.reserve(number_of_nodes);
    visited.push_back(0);
    for (auto i = 1u;  i < number_of_nodes;  ++i)
        unvisited.insert(i);
    result.reserve(number_of_nodes - 1);

    while (!unvisited.empty()) {
        Edge the_shortest = { -1, -1, INT_MAX };
        //put all edges (with their lengths) from nodes that are in visited
        for (auto node: visited) {
            auto const& lengths = graph[node];
            for (auto sec_node = 0u;  sec_node < lengths.size();  ++sec_node) {
                auto const length = lengths[sec_node];
                if (length > 0 && length < the_shortest.length
                               && unvisited.count(sec_node)) {
                    // our new closest node
                    the_shortest = {node, (int)sec_node, length};
                }
            }
        }

        //add the shortest path to the result
        result.push_back(the_shortest);

        // mark the destination as visited
        unvisited.erase(the_shortest.b);
        visited.push_back(the_shortest.b);
    };
    return result;
}

#include <iostream>

int main()
{
    /*
       2    3
   (0)--(1)--(2)
    |   / \   |
   6| 8/   \5 |7
    | /     \ |
   (3)-------(4)
         9          */
    std::vector<std::vector<int>> graph =
                     {{{0, 2, 0, 6, 0},
                       {2, 0, 3, 8, 5},
                       {0, 3, 0, 0, 7},
                       {6, 8, 0, 0, 9},
                       {0, 5, 7, 9, 0},
    }};

    auto result = minimum_spanning_tree(graph);
    for (auto const& i: result) {
        std::cout << "EDGE [ " << i.a << ", " << i.b << "], length: " << i.length << '\n';
    }
}

Answer 3

You have not used a priority queue, so performance is not good.