Advanced Dijkstra — State-Augmented Shortest Paths and K-Stop Variants [LC 787, LC 1928, Google, Amazon]

Problem Statement#

Given a directed weighted graph with non-negative weights, classic Dijkstra computes shortest distances from a single source. Advanced Dijkstra generalises the algorithm by augmenting the state with extra dimensions: number of stops used, fuel remaining, edges of a particular colour traversed, and so on. The key transformation: instead of dist[v], maintain dist[v][state] and run Dijkstra on the expanded state graph.

Constraints:

• 1 <= n <= 10^4
• 0 <= edges <= 10^5
• Weights non-negative; state dimension typically <= 100.

Example (LC 787 — cheapest flights with at most K stops):

Input: n = 4, flights = [[0,1,100],[1,2,100],[2,0,100],[1,3,600],[2,3,200]],
       src = 0, dst = 3, K = 1
Output: 700
Explanation: 0 -> 1 -> 3 = 100 + 600 = 700 (1 stop used).
             0 -> 1 -> 2 -> 3 = 400 but uses 2 stops (not allowed).

Why This Problem Matters#

Plain Dijkstra solves the basic shortest-path problem; advanced Dijkstra solves the interesting version of every shortest-path problem. The moment you add a constraint — "at most K stops," "with budget B," "while alternating between road types," "using at most one teleport" — plain shortest-path no longer applies, but state-augmented Dijkstra still does.

LeetCode 787 (Cheapest Flights Within K Stops), 1928 (Minimum Cost to Reach Destination in Time), 1976 (Number of Ways to Arrive at Destination), 1631 (Path with Minimum Effort), and 1786 (Number of Restricted Paths) all reduce to advanced Dijkstra. FAANG interviewers like these problems because they reward the candidate who recognises the underlying graph and chooses the right state expansion. The technique transfers directly to robotics planning, traffic routing with vehicle constraints, and battery-aware drone navigation.

The metaphor: classic Dijkstra walks vertices, advanced Dijkstra walks (vertex, state) pairs. As long as state transitions are deterministic and edge costs are non-negative, the standard min-heap argument still gives optimal solutions.

The Core Insight#

State-augmented Dijkstra extends the graph implicitly. Each vertex becomes a family of nodes parameterised by the state dimension. For LC 787, the state is "stops used so far": dist[v][k] is the cheapest way to reach v using exactly k stops. The min-heap pops (cost, vertex, stops) triples, and an entry is finalised once popped (under non-negative weights).

When relaxing an edge (u, v, w) from state (u, k), the resulting state is (v, k + 1) with cost dist[u][k] + w. Push only if k + 1 <= K + 1 (the +1 accounts for stops vs. flights).

For two-cost problems like LC 1928 (minimise cost while time <= maxTime), the state is time and the cost dimension is cost. Heap entries: (cost, time, vertex). We finalise based on cost; we prune based on time.

For LC 1631 (path with minimum effort, where path cost is max of edge weights, not sum), modify relaxation: new_effort = max(curr_effort, edge_weight). The min-heap pops the path with smallest maximum.

For LC 1976 (number of ways to arrive at destination using shortest paths), augment Dijkstra to maintain a count: ways[v] = sum of ways[u] over predecessors that produce the optimal distance.

The unifying pattern: classify the constraint, decide the state, augment the relaxation, run Dijkstra. The asymptotic cost is O((V * S + E * S) log (V * S)) where S is the state-space size.

Visual Dry Run (LC 787)#

n = 4, K = 1. Flights: 0->1 100, 1->2 100, 2->0 100, 1->3 600, 2->3 200. src=0, dst=3.

Min-heap starts: [(0, 0, 0)] meaning (cost, vertex, stops).

Note: stops counter equals number of intermediate cities; allow up to K + 1 = 2 edges.

Solution (Optimal)#

Python#

import heapq
from collections import defaultdict
 
def cheapest_flight_k_stops(n, flights, src, dst, K):
    """LC 787: Dijkstra augmented with stops dimension."""
    adj = defaultdict(list)
    for u, v, w in flights:
        adj[u].append((v, w))
    INF = float('inf')
    # best[v][k] = cheapest cost to reach v using k flights
    best = [[INF] * (K + 2) for _ in range(n)]
    best[src][0] = 0
 
    heap = [(0, src, 0)]  # (cost, node, flights_used)
    while heap:
        cost, u, k = heapq.heappop(heap)
        if u == dst:
            return cost
        if k > K:
            continue
        for v, w in adj[u]:
            new_cost = cost + w
            if new_cost < best[v][k + 1]:
                best[v][k + 1] = new_cost
                heapq.heappush(heap, (new_cost, v, k + 1))
    return -1
 
 
def min_cost_within_time(n, edges, max_time, fees):
    """LC 1928: minimise cost while keeping cumulative time <= max_time."""
    adj = defaultdict(list)
    for u, v, t in edges:
        adj[u].append((v, t))
        adj[v].append((u, t))
    INF = float('inf')
    # best_time_at_node[v] = minimum time we have ever reached v with cost <= current
    best_time = [INF] * n
    heap = [(fees[0], 0, 0)]  # (cost, time, node)
    while heap:
        cost, t, u = heapq.heappop(heap)
        if u == n - 1:
            return cost
        if t >= best_time[u]:
            continue
        best_time[u] = t
        for v, dt in adj[u]:
            nt = t + dt
            if nt <= max_time:
                heapq.heappush(heap, (cost + fees[v], nt, v))
    return -1

JavaScript#

class MinHeap {
    constructor() { this.h = []; }
    push(v) { this.h.push(v); this._up(this.h.length - 1); }
    pop() { const t = this.h[0], l = this.h.pop(); if (this.h.length) { this.h[0] = l; this._down(0); } return t; }
    size() { return this.h.length; }
    _cmp(a, b) { return a[0] - b[0]; }
    _up(i) { while (i > 0) { const p = (i-1)>>1; if (this._cmp(this.h[p], this.h[i]) <= 0) break; [this.h[p], this.h[i]] = [this.h[i], this.h[p]]; i = p; } }
    _down(i) { const n = this.h.length; while (true) { const l = 2*i+1, r = 2*i+2; let s = i; if (l<n && this._cmp(this.h[l], this.h[s]) < 0) s=l; if (r<n && this._cmp(this.h[r], this.h[s]) < 0) s=r; if (s===i) break; [this.h[s], this.h[i]] = [this.h[i], this.h[s]]; i = s; } }
}
 
function cheapestFlightKStops(n, flights, src, dst, K) {
    const adj = Array.from({ length: n }, () => []);
    for (const [u, v, w] of flights) adj[u].push([v, w]);
    const INF = Infinity;
    const best = Array.from({ length: n }, () => new Array(K + 2).fill(INF));
    best[src][0] = 0;
 
    const heap = new MinHeap();
    heap.push([0, src, 0]);
    while (heap.size()) {
        const [cost, u, k] = heap.pop();
        if (u === dst) return cost;
        if (k > K) continue;
        for (const [v, w] of adj[u]) {
            const nc = cost + w;
            if (nc < best[v][k + 1]) {
                best[v][k + 1] = nc;
                heap.push([nc, v, k + 1]);
            }
        }
    }
    return -1;
}

Complexity: Time O((V * S + E * S) log (V * S)) where S is the state-space dimension (e.g., K + 1 stops, T + 1 time buckets). Space O(V * S) for the augmented distance table.

Common Mistakes#

1. Forgetting to dimension dist. A 1D dist[v] cannot capture different costs at different states. Always promote to dist[v][state] when state matters.
2. Pushing stale entries without a check. Always compare against the current best for (v, state) before pushing.
3. Using BFS for K-stops. Plain BFS does not track edge weights, only edge counts. Use Dijkstra augmented with a stops counter.
4. Pruning too aggressively. When the state changes the relaxation rule (e.g., min-effort uses max instead of +), reusing classic Dijkstra rules silently breaks correctness.
5. Mixing two cost dimensions in the heap key. When optimising cost subject to time, the heap must be keyed only on cost; time is a state dimension to prune against.
6. Negative edge weights. Dijkstra's correctness rests on non-negativity. If the graph has negatives, switch to Bellman-Ford.

Interview Tips#

• Open by identifying the constraint dimension. Say: "I will augment Dijkstra with stops/time/effort as a new state dimension."
• Justify using Dijkstra rather than BFS or Bellman-Ford in 30 seconds: weights non-negative, state space finite, optimisation criterion compatible with min-heap.
• For LC 787 specifically, mention that BFS variant works because there are at most K + 1 edges; but the priority-queue version handles general weights.
• For LC 1631 (min-effort), highlight that the relaxation uses max instead of + and the heap is still keyed on the path's max edge.
• For LC 1976 (number of ways), explain that you maintain a ways[v] count alongside distance; when finding a strictly cheaper path, reset; when finding equal, accumulate.

Follow-up Questions#

1. LC 787 (Cheapest Flights Within K Stops): Use the (vertex, stops) state.
2. LC 1928 (Minimum Cost to Reach Destination in Time): Use the (vertex, time) state.
3. LC 1631 (Path with Minimum Effort): Modify relaxation to max instead of +.
4. LC 1976 (Number of Ways to Arrive at Destination): Maintain ways[] alongside dist[]; reset when strictly better, accumulate when equal.
5. What is k-th shortest path? Use Yen's algorithm or maintain a heap with up to k entries per vertex.

Key Takeaways#

• Advanced Dijkstra augments the state with extra dimensions: stops, time, fuel, colour, etc.
• The min-heap key remains the optimisation criterion; the state dimension is for pruning and filtering.
• For min-max path problems (LC 1631), replace + with max in relaxation.
• For counting problems (LC 1976), maintain ways[v] alongside distance; reset on strictly cheaper, accumulate on equal.
• Dijkstra requires non-negative weights; with negatives, switch to Bellman-Ford or SPFA.
• FAANG cue: any "shortest path with one extra constraint" calls for state-augmented Dijkstra. Identify the state, then write classic Dijkstra over the (vertex, state) pairs.