All-pairs shortest paths by repeated squaring in Python

Speeds up min-plus matrix multiplication for all-pairs shortest paths by squaring the weight matrix, needing only a logarithmic number of products.

File: apsp_repeated_squaring.py, 43 lines
Level: I0 not checked yet
Toolchain: Python 3.14, pytest 9
Test framework: pytest
Run the tests: python3 -m pytest

apsp_repeated_squaring.pyRaw GitHub

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/apsp-repeated-squaring/. The authors are listed in prov/TRREJZ.jsonl.

import math

INF = math.inf


def apsp_repeated_squaring(
    n: int, edges: list[tuple[int, int, float]]
) -> list[list[float]] | None:
    """All-pairs shortest-path weights by repeated min-plus squaring.

    The graph has the vertices 0 .. n-1 and the directed edges (u, v, w); weights may
    be negative, and parallel edges and self-loops are allowed. Returns the distance
    matrix (math.inf for no path), or None if the graph has a cycle of negative
    weight. Time Θ(n^3 log n), extra space Θ(n^2).
    """
    L = [[INF] * n for _ in range(n)]  # @line L1
    for i in range(n):
        L[i][i] = 0
    for u, v, w in edges:  # @line L2
        if w < L[u][v]:
            L[u][v] = w
    k = 1  # @line L3
    while k < n:  # @line L4
        # L[i][i] <= 0 always, so the term t = i makes M start as a copy of L
        M = [row[:] for row in L]  # @line L5
        for i in range(n):  # @line L6
            Li, Mi = L[i], M[i]
            for t in range(n):
                lit = Li[t]
                if lit == INF:
                    continue
                Lt = L[t]
                for j in range(n):  # @line L7
                    if lit + Lt[j] < Mi[j]:
                        Mi[j] = lit + Lt[j]
        L = M  # @line L8
        k *= 2
    if any(L[i][i] < 0 for i in range(n)):  # @line L9
        return None
    return L  # @line L10

test_apsp_repeated_squaring.py 82 lines

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/apsp-repeated-squaring/. The authors are listed in prov/TRREJZ.jsonl.

import math
import random

from apsp_repeated_squaring import apsp_repeated_squaring

INF = math.inf
EDGES = [(0, 1, 3), (0, 2, 8), (1, 3, 1), (1, 2, 4), (2, 1, -2), (3, 0, 2), (3, 2, 2)]


def bellman_ford(n, edges, s):
    """Independent oracle: distances from s, or None if a negative cycle is reachable."""
    d = [INF] * n
    d[s] = 0
    for _ in range(n - 1):
        for u, v, w in edges:
            if d[u] + w < d[v]:
                d[v] = d[u] + w
    if any(d[u] + w < d[v] for u, v, w in edges):
        return None
    return d


def oracle(n, edges):
    """Distance matrix from n Bellman-Ford runs, or None if any negative cycle exists."""
    rows = [bellman_ford(n, edges, s) for s in range(n)]
    return None if any(r is None for r in rows) else rows


def test_example() -> None:
    assert apsp_repeated_squaring(4, EDGES) == [
        [0, 3, 6, 4],
        [3, 0, 3, 1],
        [1, -2, 0, -1],
        [2, 0, 2, 0],
    ]


def test_negative_cycle_of_two_edges() -> None:
    # squaring only until k >= n - 1 = 1 would miss this cycle
    assert apsp_repeated_squaring(2, [(0, 1, 1), (1, 0, -2)]) is None


def test_negative_cycle_not_reachable_from_everyone() -> None:
    # 0 -> 1 only; the cycle 1 -> 2 -> 1 has weight -1 and 3 is isolated
    assert apsp_repeated_squaring(4, [(0, 1, 4), (1, 2, -3), (2, 1, 2)]) is None


def test_negative_self_loop() -> None:
    assert apsp_repeated_squaring(1, [(0, 0, -1)]) is None
    assert apsp_repeated_squaring(1, [(0, 0, 5)]) == [[0]]


def test_unreachable_and_trivial() -> None:
    assert apsp_repeated_squaring(0, []) == []
    assert apsp_repeated_squaring(3, [(0, 1, -7)]) == [[0, -7, INF], [INF, 0, INF], [INF, INF, 0]]


def test_zero_cycle_and_parallel_edges() -> None:
    edges = [(0, 1, 2), (0, 1, 1), (1, 2, -3), (2, 0, 2)]
    assert apsp_repeated_squaring(3, edges) == [[0, 1, -2], [-1, 0, -3], [2, 3, 0]]


def test_long_path_needs_all_squarings() -> None:
    n = 9
    edges = [(i, i + 1, -1) for i in range(n - 1)]
    D = apsp_repeated_squaring(n, edges)
    assert D is not None
    assert D[0][n - 1] == -(n - 1)
    assert D[n - 1][0] == INF


def test_random_against_bellman_ford() -> None:
    rng = random.Random(1976)
    for _ in range(600):
        n = rng.randint(1, 8)
        m = rng.randint(0, 3 * n)
        low = rng.choice([0, -2, -5])
        edges = [(rng.randrange(n), rng.randrange(n), rng.randint(low, 9)) for _ in range(m)]
        assert apsp_repeated_squaring(n, edges) == oracle(n, edges)

What the next level needs

the tests did not run in this build

Read the full page about All-pairs shortest paths by repeated squaring