Dijkstra Algorithm

#include<bits/stdc++.h>
#define ll long long


// NOTE:
//      graph is BI-DIRECTED (edge has 2 directions)
// Usage:
//      Dijkstra<ll> dk;
//      dk.init(n);
//      dk.addEdge(u, v, w);
//      dk.findShortestPath(source /*, sink */);
template <typename T>
struct Dijkstra 
{
    int n;
    T *dist;
    vector<pair<int, T> > *adj;

    void init(int n_) {
        n = n_;
        dist = new T[n+1];
        adj = new vector<pair<int, T> >[n+1];
    }

    void addEdge(int u, int v, T w) {
        adj[u].push_back({v, w});
        adj[v].push_back({u, w});
    }

    void findShortestPath(int source, int sink=-1) {
        fill(dist, dist + n + 1, -1);
        dist[source] = 0;

        set<pair<T, int> > bag;
        bag.insert({0, source});
        while(bag.size()) {
            T cost = bag.begin()->first;
            int u = bag.begin()->second;
            bag.erase(bag.begin());

            if (cost != dist[u]) continue;
            if (u == sink) break;

            for (auto tmp : adj[u]) {
                int v = tmp.first;
                T w = tmp.second;

                if (dist[v] == -1 || dist[v] > cost + w) {
                    dist[v] = cost + w;
                    bag.insert({cost+w, v});
                }
            }
        }
    }
};
    

Leave a Reply