leetcodealgorithms-templates6-graphs
 
// Visit www.neon.rip for more!
 
#include <vector>
#include <unordered_map>
#include <utility>
#include <queue>
 
#include <vector>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <queue>
 
using std::vector;
using std::unordered_map;
using std::unordered_set;
using std::pair;
using std::tuple;
using std::get;
using std::make_tuple;
using std::make_pair;
using std::priority_queue;
using std::greater;
 
// Given a list of edges of a connected undirected graph,
// with nodes numbered from 1 to n,
// return a list edges making up the minimum spanning tree.
vector<pair<int, int>> mst(vector<vector<int>>& edges, int n) {
    unordered_map<int, vector<pair<int, int>>> adj;
    for (int i = 1; i < n + 1; i++) {
        adj[i] = vector<pair<int, int>>();
    }
    for (vector<int> edge : edges) {
        int n1 = edge[0], n2 = edge[1], weight = edge[2];
        adj[n1].push_back(make_pair(n2, weight));
        adj[n2].push_back(make_pair(n1, weight));
    }
 
    // Initialize the heap by choosing a single node
    // (in this case 1) and pushing all its neighbors.
    priority_queue<tuple<int,int, int>, vector<tuple<int,int, int>>, greater<tuple<int, int, int>>> minHeap; 
    for (pair<int, int> neighbor : adj[1]) {
        int node = neighbor.first, weight = neighbor.second;
        minHeap.push({weight, 1, node});
    }
    
    vector<pair<int, int>> mst;
    unordered_set<int> visit;
    visit.insert(1);
    while (visit.size() < n) {
        tuple<int, int, int> cur = minHeap.top();
        minHeap.pop();
        int w1 = get<0>(cur), n1 = get<1>(cur), n2 = get<2>(cur);
 
        if (visit.count(n2) > 0) {
            continue;
        }
        mst.push_back({n1, n2});
        visit.insert(n2);
        for (pair<int, int> p : adj[n2]) {
            int neighbor = p.first, weight = p.second;
            if (visit.count(neighbor) == 0) {
                minHeap.push({weight, n2, neighbor});
            }
        }
    }
    return mst;
}