Divide-and-conquer matrix multiplication in Python

Split each matrix into four blocks and form the product from eight recursive block products; the recurrence T(n) = 8T(n/2) + Θ(1) gives Θ(n^3), the same work as the triple loop.

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

divide_and_conquer_matrix_multiplication.pyRaw GitHub

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/divide-and-conquer-matrix-multiplication/. The authors are listed in prov/4CTYKM.jsonl.

Matrix = list[list[int]]


def divide_and_conquer_matrix_multiplication(a: Matrix, b: Matrix, p: int | None = None, mod: int | None = None) -> Matrix:
    """Return a * b by the recursive eight-product method.

    a is n x m and b is m x p, as lists of rows; pass p when m = 0. With mod
    set, the arithmetic is in Z/mod Z. The recursion splits each side at its
    ceiling half, so any shape works without padding. Cost: n*m*p
    multiplications and additions; recursion depth O(log max(n, m, p)).
    """
    n, m = len(a), len(b)
    if p is None:
        p = len(b[0]) if m > 0 else 0
    if any(len(row) != m for row in a) or any(len(row) != p for row in b):
        raise ValueError("the shapes do not match")
    c = [[0] * p for _ in range(n)]
    dc_multiply_add(a, b, c, 0, 0, 0, 0, 0, 0, n, m, p, mod)
    return c


def dc_multiply_add(a: Matrix, b: Matrix, c: Matrix, ai: int, ak: int, bk: int, bj: int,
                    ci: int, cj: int, n: int, m: int, p: int, mod: int | None) -> None:
    """C view += A view * B view. The A view starts at (ai, ak), the B view at (bk, bj), the C view at (ci, cj)."""
    if n == 0 or m == 0 or p == 0:  # @line L1
        return
    if n == 1 and m == 1 and p == 1:  # @line L2
        v = c[ci][cj] + a[ai][ak] * b[bk][bj]  # @line L3
        c[ci][cj] = v % mod if mod is not None else v
        return  # @line L4
    n1, m1, p1 = (n + 1) // 2, (m + 1) // 2, (p + 1) // 2  # @line L5
    rows = ((0, n1), (n1, n - n1))
    inner = ((0, m1), (m1, m - m1))
    cols = ((0, p1), (p1, p - p1))
    for r0, rn in rows:  # @line L6
        for s0, sn in cols:  # @line L7
            k0, kn = inner[0]
            dc_multiply_add(a, b, c, ai + r0, ak + k0, bk + k0, bj + s0, ci + r0, cj + s0, rn, kn, sn, mod)  # @line L8
            k0, kn = inner[1]
            dc_multiply_add(a, b, c, ai + r0, ak + k0, bk + k0, bj + s0, ci + r0, cj + s0, rn, kn, sn, mod)  # @line L9

test_divide_and_conquer_matrix_multiplication.py 73 lines

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/divide-and-conquer-matrix-multiplication/. The authors are listed in prov/4CTYKM.jsonl.

import random

import pytest

from divide_and_conquer_matrix_multiplication import divide_and_conquer_matrix_multiplication as mul


def schoolbook(a, b, p, mod=None):
    n, m = len(a), len(b)
    c = [[0] * p for _ in range(n)]
    for i in range(n):
        for j in range(p):
            s = 0
            for k in range(m):
                s += a[i][k] * b[k][j]
            c[i][j] = s % mod if mod is not None else s
    return c


def rand(rng, r, c, lo=-9, hi=9):
    return [[rng.randint(lo, hi) for _ in range(c)] for _ in range(r)]


def test_example() -> None:
    a = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
    b = [[9, 8, 7], [6, 5, 4], [3, 2, 1]]
    assert mul(a, b) == [[30, 24, 18], [84, 69, 54], [138, 114, 90]]


def test_edges() -> None:
    assert mul([], []) == []
    assert mul([[4]], [[-3]]) == [[-12]]
    assert mul([[], [], []], [], p=2) == [[0, 0], [0, 0], [0, 0]]
    assert mul([[1, 2, 3]], [[1], [1], [1]]) == [[6]]
    with pytest.raises(ValueError):
        mul([[1, 2]], [[1]])


def test_identity_zero() -> None:
    rng = random.Random(5)
    for n in range(0, 10):
        a = rand(rng, n, n)
        eye = [[int(i == j) for j in range(n)] for i in range(n)]
        zero = [[0] * n for _ in range(n)]
        assert mul(a, eye, p=n) == a and mul(eye, a, p=n) == a
        assert mul(zero, a, p=n) == zero


def test_random_over_z() -> None:
    rng = random.Random(31337)
    for _ in range(500):
        n, m, p = rng.randint(0, 9), rng.randint(0, 9), rng.randint(0, 9)
        a, b = rand(rng, n, m, -10**6, 10**6), rand(rng, m, p, -10**6, 10**6)
        assert mul(a, b, p=p) == schoolbook(a, b, p)


def test_random_mod_p() -> None:
    rng = random.Random(8)
    for q in (2, 3, 5, 7, 13):
        for _ in range(60):
            n, m, p = rng.randint(1, 8), rng.randint(1, 8), rng.randint(1, 8)
            a, b = rand(rng, n, m, 0, q - 1), rand(rng, m, p, 0, q - 1)
            assert mul(a, b, mod=q) == schoolbook(a, b, p, q)


def test_powers_of_two_and_odd() -> None:
    rng = random.Random(77)
    for n in (1, 2, 4, 8, 16, 3, 5, 7, 15, 17):
        a, b = rand(rng, n, n), rand(rng, n, n)
        assert mul(a, b) == schoolbook(a, b, n)

What the next level needs

the tests did not run in this build

Read the full page about Divide-and-conquer matrix multiplication