BFS and Dijkstra with Extra State

Every traversal so far has assumed a node is a place. Plenty of interview problems break that assumption: you may cross at most k obstacles, you may take at most k stops, you hold a set of keys, you have one wall-breaking charge left. Arriving at the same cell with a different budget is not the same situation, and treating it as such is the bug.

The fix is a single idea: make the state a tuple, not a node. Everything else — BFS, Dijkstra, the visited set — stays exactly as you already know it. This is the family that LC 787, 864 and 1293 belong to, and it is under-taught relative to how often it comes up.

What you’ll learn

• Why a plain visited set is wrong here, and what to key it on instead.
• How to decide whether the extra dimension needs BFS or Dijkstra.
• The bitmask trick for “collect keys” problems.
• The one-line change that turns LC 787 from wrong to right.

The cue

You need extra state when…

• The problem gives you a budget: at most k stops, k obstacles you may remove, one wall you may break, k moves of a special kind.
• You collect something that changes what is reachable: keys that open doors, items, a set of visited targets.
• Movement depends on a mode: fuel level, current direction, whose turn it is, step parity.

The test that tells you immediately: can arriving at the same node twice, in different circumstances, lead to different outcomes? If yes, the node alone is not the state, and a visited set keyed on the node will discard the arrival you needed.

The state is then a tuple — (node, budget), (cell, keys_bitmask), (cell, obstacles_used) — and every part of the algorithm keys on the whole tuple. The graph you are searching is the product of the original graph and the state space, which is why the complexity multiplies rather than adds.

Reach for plain BFS or Dijkstra when the answer depends only on where you are. Adding a dimension you do not need multiplies the work for nothing.

Visual intuition

The traversal machinery is unchanged — the same queue, the same frontier. What changes is what a queue entry contains:

graphDijkstra proper — then imagine every node duplicated per budget levelthe machinery you already know

heap

A:0

1import heapq2 3def dijkstra(graph, start):4    dist = {v: inf for v in graph}5    dist[start] = 06    pq = [(0, start)]7    while pq:8        d, v = heapq.heappop(pq)9        if d > dist[v]: continue      # stale entry — skip10        for u, wt in graph[v]:11            nd = d + wt12            if nd < dist[u]:13                dist[u] = nd14                heapq.heappush(pq, (nd, u))15    return dist

seedEvery distance starts at infinity except A, which is 0. Dijkstra's promise: whichever unsettled node has the smallest tentative distance already has its *final* distance — provided no edge weight is negative.

Play‹›1/6⟲

With a stop limit, node B is no longer one node: (B, 1 stop used) and (B, 2 stops used) are separate states with separate distances. The picture is this graph replicated once per budget value, and the search runs over the whole stack of copies.

Why the plain visited set fails

LC 1293, “shortest path in a grid with obstacle elimination”: you may remove up to k obstacles. Take k = 1 and this grid, where # is an obstacle:

text

. # .
. # .
. . .

Two routes reach the bottom-right. The direct one goes straight down column 0 and across — no obstacles, length 4. Another goes through (0,1) by spending the one removal.

Now suppose the search reaches cell (1,0) twice: once having spent 0 removals, once having spent 1. A visited set keyed on (1,0) alone accepts the first arrival and discards the second — or worse, accepts whichever arrives first and rejects a later arrival that has more budget left.

arrival at (1,0)	removals used	still reachable from here?
via (0,0)	0	everything
via (0,1)	1	fewer options — the budget is gone

These are genuinely different situations. Keying visited on (row, col, used) keeps both. Keying it on (row, col) throws away the one you needed, and the answer comes back too large or -1.

The symptom is a plausible wrong answer, not a crash

A node-keyed visited set does not error. It returns a longer path, or claims the target is unreachable, on some inputs. That makes it one of the harder bugs to spot in an interview — which is why the state design should be stated out loud before any code is written.

The template

state_space_search.py

from collections import deque
import heapq
 
 
def shortest_path_with_removals(grid, k):        # LC 1293
    rows, cols = len(grid), len(grid[0])
    if k >= rows + cols - 2:                      # enough budget to go straight
        return rows + cols - 2
 
    start = (0, 0, k)                             # (row, col, removals LEFT)
    q = deque([(start, 0)])
    seen = {start}                                # keyed on the WHOLE tuple
 
    while q:
        (r, c, left), dist = q.popleft()
        if (r, c) == (rows - 1, cols - 1):
            return dist
        for dr, dc in ((-1, 0), (1, 0), (0, -1), (0, 1)):
            nr, nc = r + dr, c + dc
            if not (0 <= nr < rows and 0 <= nc < cols):
                continue
            nleft = left - grid[nr][nc]             # spend one on an obstacle
            if nleft < 0:
                continue
            state = (nr, nc, nleft)
            if state not in seen:
                seen.add(state)
                q.append((state, dist + 1))
    return -1
 
 
def cheapest_flight_k_stops(n, flights, src, dst, k):   # LC 787
    graph = {}
    for u, v, w in flights:
        graph.setdefault(u, []).append((v, w))
 
    # Dijkstra over (cost, node, stops_used). Note best[] is per (node, stops).
    pq = [(0, src, 0)]
    best = {}
    while pq:
        cost, node, stops = heapq.heappop(pq)
        if node == dst:
            return cost
        if stops > k:
            continue
        # Only skip if we have been here before with FEWER OR EQUAL stops.
        if best.get(node, float("inf")) <= stops:
            continue
        best[node] = stops
        for nxt, w in graph.get(node, []):
            heapq.heappush(pq, (cost + w, nxt, stops + 1))
    return -1
 
 
grid = [[0, 1, 0], [0, 1, 0], [0, 0, 0]]
print(shortest_path_with_removals(grid, 1))     # expect 4
print(cheapest_flight_k_stops(
    4, [[0,1,100],[1,2,100],[2,0,100],[1,3,600],[2,3,200]], 0, 3, 1))  # expect 700

Dry run

LC 787, n = 4, flights 0→1 (100), 1→2 (100), 2→0 (100), 1→3 (600), 2→3 (200), from 0 to 3 with at most k = 1 stop.

pop	(cost, node, stops)	action
1	(0, 0, 0)	push (100, 1, 1)
2	(100, 1, 1)	stops = 1 ≤ k, so push (200, 2, 2) and (700, 3, 2)
3	(200, 2, 2)	stops = 2 > k — discard, cannot extend
4	(700, 3, 2)	node is the destination → return 700

The cheapest route by cost is 0→1→2→3 at 500, but it uses two stops and is therefore illegal. Plain Dijkstra returns 500 and is wrong.

Two things worth reading off that table:

• The stop check happens on pop, not on push. Popping (200, 2, 2) and discarding it is correct: an over-budget state may still be reached, it just may not be extended.
• Node 3 is reached at cost 700 with 2 stops — which is legal, because the limit counts intermediate stops, and the destination is not one. Off-by-one here is the second most common bug on LC 787, after ignoring stops entirely.

LC 787 is the problem where “just use Dijkstra” is a trap

Dijkstra’s correctness rests on “once a node is settled, its distance is final”. With a stop limit that is false: node 2 might be settled cheaply via a long route that has no budget left, blocking the expensive-but-legal route. Either key the settled-set on (node, stops), or use Bellman-Ford relaxed exactly k + 1 times, which is the cleaner answer and worth naming.

Complexity

Problem	State space	Time
LC 1293 (obstacles, BFS)	$rc \cdot (k+1)$	$O(rc \cdot k)$
LC 787 (stops, Bellman-Ford)	$n \cdot (k+1)$	$O(k \cdot E)$
LC 864 (keys, BFS + bitmask)	$rc \cdot 2^{\text{keys}}$	$O(rc \cdot 2^{\text{keys}})$
LC 1928 (fuel, Dijkstra)	$n \cdot \text{maxFuel}$	$O(n \cdot f \log(nf))$

The rule: multiply, do not add. The state space is the node count times the number of distinct budget values, and the runtime follows. That is also the constraint check — LC 864 caps keys at 6 precisely because $2^6 = 64$ keeps $rc \cdot 64$ tractable, and seeing a small cap like that in a problem statement is a strong hint that a bitmask dimension is intended.

The variant map

Extra state	Encoding	Search	Canonical problem
Budget of removals	(r, c, left)	BFS — all moves cost 1	1293 Obstacle Elimination
Stop / edge limit	(node, stops)	Bellman-Ford, or Dijkstra keyed on the tuple	787 Cheapest Flights Within K Stops
Collected keys	(r, c, bitmask)	BFS	864 Shortest Path to Get All Keys
Fuel remaining	(node, fuel)	Dijkstra	1928 Minimum Cost With Deadline
Step parity / turn	(node, step % m)	BFS	1654 Minimum Jumps to Reach Home
Visited-target set	(node, bitmask)	BFS	847 Shortest Path Visiting All Nodes

Bitmask for collected things, and why it works

When you collect from a small fixed set — up to 6 keys, up to 12 target nodes — encode the collection as an integer where bit i means “have item i”. Then (cell, mask) is a plain tuple of integers: hashable, cheap to compare, and mask | (1 << i) adds an item in one operation.

LC 847 is the same idea inverted: the mask records which nodes have been visited, and the answer is the first state whose mask is all-ones. Recognising that “visit all of a small set” and “collect all of a small set” are the same problem is the transferable part.

Pitfalls

• Keying visited on the node instead of the full state. The defining bug. Returns a plausible wrong answer on some inputs and no error on any.
• Checking the budget on push rather than on pop in LC 787. An over-budget state may be reached; it may not be extended.
• Off-by-one on “stops”. k stops means k + 1 edges. Read the statement twice.
• Using plain Dijkstra with a stop limit. Its settling argument does not hold when a cheap arrival can be budget-exhausted. Key on the tuple, or use Bellman-Ford.
• Adding a dimension you do not need. If the outcome depends only on position, the extra dimension multiplies the work for nothing.
• Forgetting the early-exit shortcut. In LC 1293, k >= rows + cols - 2 means you can walk straight through everything; without that check large k explodes the state space needlessly.

Interview follow-ups

They ask	What they’re checking	The answer
“Why not a normal visited set?”	The core insight	Because the same node reached with a different budget is a different situation. Key visited on the whole tuple, or you discard the arrival you needed
“BFS or Dijkstra here?”	Whether you check edge weights	BFS if every move costs the same, Dijkstra if costs differ. The extra dimension does not change that choice
“What is the state space size?”	Complexity reasoning	Nodes times budget values. That multiplication is both the runtime and the reason problems cap keys at 6
“Plain Dijkstra gives 500 on LC 787. Why is that wrong?”	Whether you know the failure	500 uses two stops and the limit is one. Dijkstra’s settling argument fails because a cheap arrival may have no budget left
“Give a cleaner solution to LC 787”	Breadth	Bellman-Ford relaxed exactly k + 1 times — each round extends every path by one edge, so the round count is the edge budget
“How would you encode 6 collectable keys?”	Practical modelling	A 6-bit mask; state is (r, c, mask) and mask | (1 << i) collects key i. Goal is mask == (1 << n) - 1
“The budget is 10^9”	Whether you spot the degenerate case	Then it does not constrain you — fall back to the unconstrained search. Any budget above the maximum useful path length is equivalent to infinite

Practice

3 problems

0 easy0 medium3 hard

Work down the ladder. Tick each problem off as you go — progress is saved in this browser, and the Export button in the filter bar writes it to a file you can keep.

• ○847Shortest Path Visiting All Nodeshard
• ○864Shortest Path to Get All Keyshard
• ○1293Shortest Path in a Grid with Obstacles Eliminationhard

Exercises

LC 1293 — Shortest Path with Obstacle Elimination · Hard

pch.exerciseTag pch.exerciseTitle

LC 787 — Cheapest Flights Within K Stops · Medium

pch.exerciseTag pch.exerciseTitle

LC 864 — Shortest Path to Get All Keys · Hard (bitmask state)

pch.exerciseTag pch.exerciseTitle

Self-check

pch.quizTag Search with extra state — self-check

1. What is the test for whether you need extra state in the search?

   A Whether the graph is weighted B Whether arriving at the same node in different circumstances can lead to different outcomes — if yes, the node alone is not the state C Whether the graph has cycles D Whether the grid is larger than 100x100
   pch.quizShowAnswer

   B — Whether arriving at the same node in different circumstances can lead to different outcomes — if yes, the node alone is not the state — That one question decides it. If a budget, a key set or a mode changes what is reachable from here, then (node) is not a state and a node-keyed visited set will discard the arrival you needed.

2. A node-keyed visited set on LC 1293 produces what symptom?

   A An infinite loop B A plausible wrong answer on some inputs — a longer path or a false -1 — with no error at all C A KeyError D Correct answers, just slower
   pch.quizShowAnswer

   B — A plausible wrong answer on some inputs — a longer path or a false -1 — with no error at all — This is what makes it dangerous. It never crashes, so it survives casual testing. That is why the state design should be stated out loud before any code is written.

3. Plain Dijkstra returns 500 on LC 787 when the answer is 700. Why?

   A A bug in the heap comparison B The 500 route uses two stops and the limit is one — Dijkstra's 'once settled, final' argument fails because a cheap arrival may have no budget left C Because the graph has negative weights D Because 500 is correct and 700 is wrong
   pch.quizShowAnswer

   B — The 500 route uses two stops and the limit is one — Dijkstra's 'once settled, final' argument fails because a cheap arrival may have no budget left — The settling argument is the foundation of Dijkstra's correctness, and a budget dimension invalidates it. Either key the settled set on (node, stops), or use Bellman-Ford with k+1 rounds.

4. In the Bellman-Ford solution to LC 787, why snapshot the distance array each round?

   A For efficiency B So that a single round cannot chain two edges — one round must extend every path by exactly one edge, since the round count IS the edge budget C To handle negative cycles D It is not necessary
   pch.quizShowAnswer

   B — So that a single round cannot chain two edges — one round must extend every path by exactly one edge, since the round count IS the edge budget — Without the snapshot, one round can relax u then immediately relax v using the freshly-updated u, using two edges in one round. The stop limit then stops meaning anything.

5. How do you size the state space, and what does it tell you?

   A Nodes plus budget values B Nodes TIMES budget values — which is also why LC 864 caps keys at 6, since 2^6 = 64 keeps rc·64 tractable C Just the node count D It cannot be bounded
   pch.quizShowAnswer

   B — Nodes TIMES budget values — which is also why LC 864 caps keys at 6, since 2^6 = 64 keeps rc·64 tractable — Multiply, do not add. And read it backwards: a suspiciously small cap in the constraints (keys <= 6, nodes <= 12) is a strong hint that a bitmask dimension is intended.

Recall card

• Cue — a budget (k stops, k removals), a collection (keys, targets), or a mode (fuel, parity). Test: can the same node in different circumstances give different outcomes?
• The fix — state is a tuple, not a node: (node, budget), (cell, mask). Key the visited set on the whole tuple.
• BFS or Dijkstra — unchanged rule: equal move costs → BFS, differing → Dijkstra. The extra dimension does not affect that choice.
• State space multiplies — nodes × budget values. That is the runtime and the reason caps are small.
• Small caps are a hint — keys ≤ 6, nodes ≤ 12 means a bitmask dimension.
• LC 787 — plain Dijkstra is wrong. Bellman-Ford with k + 1 rounds and a per-round snapshot is the clean answer.

Recap

• A node is a place; a state is a place plus everything that changes what you can do from it. Most of this family is just noticing the difference.
• The visited set must key on the full state. Getting that wrong gives wrong answers silently rather than crashing.
• The state space is the product of the graph and the budget, so the runtime multiplies — and an unusually small constraint is the problem telling you a bitmask is intended.
• LC 787 is the cautionary case: the shortest path by cost can be illegal, and Dijkstra’s settling argument does not survive a budget.

Next: Union-Find Problem Patterns — connectivity without a traversal.