reachable-nodes-in-subdivided-graph.cpp

// Time:  O((|E| + |V|) * log|V|) = O(|E| * log|V|),
//        if we can further to use Fibonacci heap, it would be O(|E| + |V| * log|V|)
// Space: O(|E| + |V|) = O(|E|)

class Solution {
public:
    int reachableNodes(vector<vector<int>>& edges, int M, int N) {
        using P = pair<int, int>;
        vector<vector<P>> adj(N);
        for (const auto& edge: edges) {
            int u = edge[0], v = edge[1], w = edge[2];
            adj[u].emplace_back(v, w);
            adj[v].emplace_back(u, w);
        }
        unordered_map<int, int> best;
        best[0] = 0;
        unordered_map<int, unordered_map<int, int>> count;
        int result = 0;
        priority_queue<P, vector<P>, greater<P>> min_heap;
        min_heap.emplace(0, 0);
        while (!min_heap.empty()) {
            int curr_total, u;
            tie(curr_total, u) = min_heap.top(); min_heap.pop();
            if (best.count(u) && best[u] < curr_total) {
                continue;
            }
            ++result;
            for (const auto& kvp: adj[u]) {
                int v, w;
                tie(v, w) = kvp;
                count[u][v] = min(w, M - curr_total);
                int next_total = curr_total + w + 1;
                if (next_total <= M && 
                    (!best.count(v) || next_total < best[v])) {
                    best[v] = next_total;
                    min_heap.emplace(next_total, v);
                }
            }
        }
        for (const auto& edge: edges) {
            int u = edge[0], v = edge[1], w = edge[2];
            result += min(w, count[u][v] + count[v][u]);
        }
        return result;
    }
};