Slow all-pairs shortest paths by matrix multiplication in Python

Computes all-pairs shortest paths by extending path lengths one edge at a time with a min-plus matrix product, in cubic time per step.

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

slow_apsp_matrix_multiplication.pyRaw GitHub

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/slow-apsp-matrix-multiplication/. The authors are listed in prov/EK3RJN.jsonl.

import math

INF = math.inf


def weight_matrix(n: int, edges: list[tuple[int, int, float]]) -> list[list[float]]:
    """W[i][i] = min(0, self-loops), W[i][j] = least weight of an arc i -> j, INF if none."""
    w = [[0 if i == j else INF for j in range(n)] for i in range(n)]
    for u, v, c in edges:
        if c < w[u][v]:
            w[u][v] = c
    return w


def extend(l: list[list[float]], w: list[list[float]]) -> list[list[float]]:
    """One min-plus product: walks of at most r arcs times one more arc (or none)."""
    n = len(l)
    out = [[INF] * n for _ in range(n)]  # @line L6
    for i in range(n):  # @line L7
        li, oi = l[i], out[i]
        for j in range(n):
            best = INF
            for k in range(n):  # @line L8
                cand = li[k] + w[k][j]
                if cand < best:
                    best = cand
            oi[j] = best
    return out  # @line L9


def slow_apsp(n: int, edges: list[tuple[int, int, float]]) -> list[list[float]] | None:
    """All-pairs distances by n - 1 min-plus products; INF marks an unreachable pair.

    Vertices are 0 .. n-1 and edges are (u, v, weight) with any real weights.
    Returns None when the graph has a negative cycle. Time Θ(n^4), space Θ(n^2).
    """
    w = weight_matrix(n, edges)
    l = [row[:] for row in w]  # @line L1: walks with at most 1 arc
    for _ in range(2, n + 1):  # @line L2
        l = extend(l, w)  # @line L3
    if any(l[i][i] < 0 for i in range(n)):  # @line L4: n arcs are enough to close any cycle
        return None
    return l  # @line L5

test_slow_apsp_matrix_multiplication.py 76 lines

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/slow-apsp-matrix-multiplication/. The authors are listed in prov/EK3RJN.jsonl.

import math
import random
import unittest

from slow_apsp_matrix_multiplication import extend, slow_apsp, weight_matrix

INF = math.inf
CLRS = [(0, 1, 3), (0, 2, 8), (0, 4, -4), (1, 3, 1), (1, 4, 7), (2, 1, 4), (3, 0, 2), (3, 2, -5), (4, 3, 6)]


def bellman_ford_all(n, edges):
    """Independent oracle: Bellman-Ford from every source, with a final pass for negative cycles."""
    rows = []
    for s in range(n):
        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
        for u, v, w in edges:
            if d[u] + w < d[v]:
                return None
        rows.append(d)
    return rows


class TestSlowAPSP(unittest.TestCase):
    def test_clrs_example(self):
        self.assertEqual(
            slow_apsp(5, CLRS),
            [[0, 1, -3, 2, -4], [3, 0, -4, 1, -1], [7, 4, 0, 5, 3], [2, -1, -5, 0, -2], [8, 5, 1, 6, 0]],
        )

    def test_intermediate_matrix(self):
        w = weight_matrix(5, CLRS)
        self.assertEqual(extend(w, w)[0], [0, 3, 8, 2, -4])
        self.assertEqual(extend(w, w)[4], [8, INF, 1, 6, 0])

    def test_negative_cycles(self):
        self.assertIsNone(slow_apsp(5, CLRS + [(2, 3, 4)]))
        self.assertIsNone(slow_apsp(1, [(0, 0, -1)]))
        # a negative cycle through all n vertices needs the n-th product
        self.assertIsNone(slow_apsp(3, [(0, 1, 1), (1, 2, 1), (2, 0, -3)]))
        # a negative cycle that most vertices cannot reach is still reported
        self.assertIsNone(slow_apsp(4, [(0, 1, 1), (2, 3, -1), (3, 2, -1)]))

    def test_unreachable_and_small(self):
        self.assertEqual(slow_apsp(0, []), [])
        self.assertEqual(slow_apsp(1, [(0, 0, 5)]), [[0]])
        self.assertEqual(slow_apsp(3, [(0, 1, -2), (0, 1, 4)]), [[0, -2, INF], [INF, 0, INF], [INF, INF, 0]])

    def test_zero_cycle(self):
        self.assertEqual(slow_apsp(2, [(0, 1, 3), (1, 0, -3)]), [[0, 3], [-3, 0]])

    def test_random_against_bellman_ford(self):
        rng = random.Random(1974)
        negative = 0
        for _ in range(600):
            n = rng.randint(0, 6)
            m = rng.randint(0, 3 * n)
            low = rng.choice([0, -2, -4])
            edges = [(rng.randrange(n), rng.randrange(n), rng.randint(low, 9)) for _ in range(m)] if n else []
            expected = bellman_ford_all(n, edges)
            got = slow_apsp(n, edges)
            if expected is None:
                negative += 1
            self.assertEqual(got, expected)
        self.assertGreater(negative, 50)


if __name__ == "__main__":
    unittest.main()

What the next level needs

the tests did not run in this build

Read the full page about Slow all-pairs shortest paths by matrix multiplication