Prim's Algorithm

#include <iostream>
#include <vector>
#include <queue>

#define INF 1e9

using namespace std;

pair<int, vector<int>> prim(vector<vector<pair<int, int>>> &graph, int start) {
    // Min heap storing { weight, node }
    priority_queue<pair<int, int>, vector<pair<int, int>>, greater<pair<int, int>>> queue;
    // We start with the start node. It has an edge weight of 0, obviously!
    queue.push({ 0, start });

    // We keep track of the nodes we've included in the MST
    vector<int> visited(graph.size());

    // We keep track of the minimum edge we've found for each node so far
    vector<int> min_edge(graph.size(), INF);

    // We keep track of the node connected using minimum edge for each node
    vector<int> parents(graph.size(), -1);

    // When we add a node, we include it's edge weight
    int total_mst_cost = 0;

    // We keep going until we exhaust all the nodes.
    while (!queue.empty()) {
        // destructuring pair<weight, node>
        auto [ weight, node ] = queue.top();
        queue.pop();

        // We already added the node into MST, ignore it
        if (visited[node]) continue;

        // We include this node into the tree, add it's edge weight
        visited[node] = true;
        total_mst_cost += weight;

        // Iterate through each of the edges and check if any of the edges have smaller weight than what we already discovered.
        for (auto &edge: graph[node]) {
            auto [ neighbor, edge_weight ] = edge;

            // If we already included the node into the MST, we ignore this
            if (!visited[neighbor]) {
                if (edge_weight < min_edge[neighbor]) {
                    // We do not include the node into the MST right here,
                    // because we can potentially find a better edge weight. For example,
                    // if we find an edge weight { 10, B }, we push it into the queue. After that, we may find { 2, B }. The priority queue will put this before the { 10, B } edge. So, this will cause the better edge to be calculated, turning visited[B] = true, and dropping { 10, B } edge.
                    parents[neighbor] = node;
                    min_edge[neighbor] = edge_weight;
                    queue.push({ edge_weight, neighbor });
                }
            }
        }
    }

    return { total_mst_cost, parents };
}

int main() {
    int nodes, edges, start;
    cin >> nodes >> edges >> start;

    vector<vector<pair<int, int>>> graph(nodes + 1);

    for (int i = 0; i < edges; i++) {
        int from, to, weight;
        cin >> from >> to >> weight;
        graph[from].push_back({ to, weight });
        graph[to].push_back({ from, weight });
    }

    auto [ total_cost, parents ] = prim(graph, start);

    cout << "Total Cost of Minimum Spanning Tree: " << total_cost << "\n";

    // Print the edges that make up the tree
    cout << "Edges used in the MST:\n";
    for (int i = 1; i <= nodes; i++) { // Start at 2 because 1 is the root
        if (i != start && parents[i] != -1) {
            cout << parents[i] << " - " << i << "\n";
        }
    }
}