cswork / libraryOpen practice workspace

Algorithm tutorials

Shortest Paths: Dijkstra, Bellman-Ford, Floyd-Warshall

How to pick and implement a shortest-path algorithm: heap-based Dijkstra for non-negative weights, Bellman-Ford for negative edges and edge-count limits, Floyd-Warshall for all pairs on small graphs.

Intuition

Every shortest-path algorithm is built from one operation, relaxation: if going to u and then taking edge (u, v, w) beats the best known distance to v, update it, i.e. dist[v] = min(dist[v], dist[u] + w).

The algorithms differ only in which edges they relax and in what order:

  • Dijkstra (non-negative weights) always settles the unsettled vertex with the smallest tentative distance. Any other route to it would leave the settled region through a vertex that is at least as far away, and non-negative edges can only add to that, so the distance is final the moment it is popped.
  • Bellman-Ford relaxes every edge, round after round. After round i, every shortest path that uses at most i edges is correct. Simple paths have at most n - 1 edges, so n - 1 rounds suffice; if round n still improves something, a negative cycle is reachable.
  • Floyd-Warshall answers all pairs. After processing intermediate vertex k, dist[i][j] is the best path whose inner vertices all come from {0..k}.

When to use it in interviews / OAs

SituationAlgorithmTime
Unweighted edgesBFSO(n + m)
Weights only 0 or 10-1 BFSO(n + m)
Non-negative weights, one sourceDijkstra with a heapO(m log n)
Negative edges, or "at most k edges"Bellman-FordO(n m), or O(k m)
All pairs, n up to a few hundredFloyd-WarshallO(n^3)
Graph is a DAGDP in topological orderO(n + m)

Signals: "minimum cost / time / distance between", "network delay", "cheapest route", "within k stops", "for every pair of cities", "threshold distance".

Template

Dijkstra with a binary heap

import heapq

INF = float("inf")


def dijkstra(adj: list[list[tuple[int, int]]], src: int) -> tuple[list[float], list[int]]:
    n = len(adj)
    dist = [INF] * n
    parent = [-1] * n
    dist[src] = 0
    pq = [(0, src)]
    while pq:
        d, u = heapq.heappop(pq)
        if d > dist[u]:  # stale entry: u was already settled with a smaller distance
            continue
        for v, w in adj[u]:
            nd = d + w
            if nd < dist[v]:
                dist[v] = nd
                parent[v] = u
                heapq.heappush(pq, (nd, v))
    return dist, parent

Python's heapq has no decrease-key, so we push duplicates and skip stale ones on pop ("lazy deletion"). To rebuild a path, follow parent from the target back to src.

Counting shortest paths

A frequent follow-up: how many different shortest routes reach the destination? Keep a counter next to the distance: a strictly better distance resets it, an equal distance adds to it.

def count_shortest_paths(n: int, edges: list[tuple[int, int, int]], src: int, dst: int,
                         mod: int = 10**9 + 7) -> int:
    adj: list[list[tuple[int, int]]] = [[] for _ in range(n)]
    for u, v, w in edges:  # undirected roads, positive times
        adj[u].append((v, w))
        adj[v].append((u, w))
    dist = [INF] * n
    ways = [0] * n
    dist[src], ways[src] = 0, 1
    pq = [(0, src)]
    while pq:
        d, u = heapq.heappop(pq)
        if d > dist[u]:
            continue
        for v, w in adj[u]:
            nd = d + w
            if nd < dist[v]:
                dist[v], ways[v] = nd, ways[u]
                heapq.heappush(pq, (nd, v))
            elif nd == dist[v]:
                ways[v] = (ways[v] + ways[u]) % mod
    return ways[dst]

This is correct because u is final when popped, so every contribution it passes on is complete (all weights must be positive).

Bellman-Ford with negative-cycle detection

def bellman_ford(n: int, edges: list[tuple[int, int, int]], src: int) -> list[float] | None:
    dist = [INF] * n
    dist[src] = 0
    for _ in range(n - 1):
        changed = False
        for u, v, w in edges:
            if dist[u] + w < dist[v]:
                dist[v] = dist[u] + w
                changed = True
        if not changed:  # converged early
            break
    for u, v, w in edges:  # still improvable: a negative cycle is reachable from src
        if dist[u] + w < dist[v]:
            return None
    return dist

At most k edges ("cheapest flights within k stops")

Run exactly k rounds, and relax from a copy of the previous round. Otherwise one round can chain several edges and break the edge budget.

def cheapest_within_k_edges(n: int, edges: list[tuple[int, int, int]],
                            src: int, dst: int, k: int) -> int:
    dist = [INF] * n
    dist[src] = 0
    for _ in range(k):
        prev = dist[:]
        for u, v, w in edges:
            if prev[u] + w < dist[v]:
                dist[v] = prev[u] + w
    return -1 if dist[dst] == INF else dist[dst]

"Within K stops" means at most K + 1 flights, so call it with k = K + 1.

Floyd-Warshall

def floyd_warshall(n: int, edges: list[tuple[int, int, int]]) -> list[list[float]]:
    dist = [[INF] * n for _ in range(n)]
    for i in range(n):
        dist[i][i] = 0
    for u, v, w in edges:
        dist[u][v] = min(dist[u][v], w)
    for k in range(n):  # k must be the OUTER loop
        dk = dist[k]
        for i in range(n):
            dik = dist[i][k]
            if dik == INF:
                continue
            di = dist[i]
            for j in range(n):
                if dik + dk[j] < di[j]:
                    di[j] = dik + dk[j]
    return dist  # dist[i][i] < 0 means i lies on a negative cycle

Complexity

  • Dijkstra with a heap: O((n + m) log n) time, O(n + m) space. On a dense graph (m close to n^2) an O(n^2) array-scan version is faster.
  • Bellman-Ford: O(n * m) worst case, O(k * m) for the k-edge variant. The queue-based variant (often called SPFA) is usually faster but has the same worst case.
  • Floyd-Warshall: O(n^3) time, O(n^2) space. About 10^8 basic steps at n = 460, which is already slow in Python.

Common pitfalls

  • Dijkstra with negative edges can settle a vertex too early and return wrong answers. Use Bellman-Ford, or a DAG DP if the graph is acyclic.
  • Not skipping stale heap entries. Without if d > dist[u]: continue, the same vertex is expanded repeatedly and the run can blow up.
  • Wrong loop order in Floyd-Warshall. The intermediate vertex k must be the outermost loop.
  • Bellman-Ford's verdict is relative to the source. "No negative cycle found" only means none is reachable from src. To test the whole graph, add a virtual source with 0-weight edges to every vertex.
  • An undirected edge with negative weight is itself a negative cycle (walk it back and forth).
  • Overflow-style infinity. In languages with fixed-width integers, INF + w overflows; check dist[u] != INF before adding. In Python, float("inf") is safe.
  • Off-by-one in "k stops". Stops are intermediate vertices, so k stops means k + 1 edges.

Interview variations

  • Network delay time. Dijkstra from the source, answer is the maximum distance (or -1 if some node is unreachable).
  • Path with maximum probability. Multiply instead of add and pop the largest first; this works because every factor is at most 1.
  • Minimum effort / swim in rising water. The path cost is the maximum edge on it; Dijkstra still works with max(d, w) in place of d + w (alternatives: binary search plus BFS, or union-find over sorted edges).
  • Number of ways to arrive at destination. The counting Dijkstra above.
  • City with the fewest neighbors within a threshold. Floyd-Warshall, then count per row.
  • Transitive closure. Floyd-Warshall with booleans (reach[i][j] |= reach[i][k] and reach[k][j]) answers "is course a a prerequisite of b" for many queries.
  • Second-shortest or k-th shortest arrival. Let each vertex be popped up to k times instead of once.
  • State-augmented Dijkstra. Discounts, fuel, keys or stop counts become part of the vertex: (node, extra).