dijkshtras_algorithm.cpp

#include "metro.hpp"

vector<int> dijk(int src, bool avoid_transfers) {
    vector<int> via(total_stations);
    priority_queue<pair<double, int>, vector<pair<double, int>>, greater<pair<double, int>>> pq;
    mp.clear();
    set<int> vis;

    for(int i = 0; i < total_stations; i++) {
        mp[i] = INF;
        via[i] = i;
    }

    // If the source itself is closed, there is nowhere to route from.
    if (station_closed[src]) {
        return via;
    }

    mp[src] = 0.0;
    pq.push({0.0, src});

    while(!pq.empty()) {
        auto curr = pq.top();
        pq.pop();

        if(vis.count(curr.second)) continue;
        vis.insert(curr.second);

        for(int i = 0; i < total_stations; i++) {
            // A closed station accepts no traffic, in either direction.
            if (station_closed[i]) continue;

            if(graph[curr.second][i] > 0.0 && !vis.count(i)) {

                double penalty = 0.0;
                if (avoid_transfers && getLine(curr.second) != getLine(i)) {
                    penalty = 15.0; // Penalty for changing lines
                }

                double newDist = curr.first + graph[curr.second][i] + penalty;
                if(mp[i] > newDist) {
                    mp[i] = newDist;
                    pq.push({newDist, i});
                    via[i] = curr.second;
                }
            }
        }
    }
    return via;
}