Strassen's algorithm in Python
Multiply 2 by 2 block matrices with seven recursive block products instead of eight, which gives n^(log2 7) multiplications and O(n^(log2 7)) ring operations, about n^2.807.
- File
strassen_algorithm.py, 94 lines- Level
- I0 not checked yet
- Toolchain
- Python 3.14, pytest 9
- Test framework
- pytest
- Run the tests
python3 -m pytest
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/strassen-algorithm/. The authors are listed in prov/MZ8PBV.jsonl.
Matrix = list[list[int]]
def _classical(a: Matrix, b: Matrix, p: int, mod: int | None) -> Matrix:
m = len(b)
c = [[0] * p for _ in range(len(a))]
for i, row in enumerate(a):
ci = c[i]
for k in range(m):
aik = row[k]
if aik:
bk = b[k]
for j in range(p):
ci[j] += aik * bk[j]
if mod is not None:
c[i] = [x % mod for x in ci]
return c
def _add(x: Matrix, y: Matrix, mod: int | None) -> Matrix:
if mod is None:
return [[u + v for u, v in zip(r, s)] for r, s in zip(x, y)]
return [[(u + v) % mod for u, v in zip(r, s)] for r, s in zip(x, y)]
def _sub(x: Matrix, y: Matrix, mod: int | None) -> Matrix:
if mod is None:
return [[u - v for u, v in zip(r, s)] for r, s in zip(x, y)]
return [[(u - v) % mod for u, v in zip(r, s)] for r, s in zip(x, y)]
def _block(x: Matrix, r0: int, r1: int, c0: int, c1: int) -> Matrix:
return [row[c0:c1] for row in x[r0:r1]]
def strassen_algorithm(a: Matrix, b: Matrix, p: int | None = None, mod: int | None = None, n0: int = 1) -> Matrix:
"""Return a * b by Strassen's seven-product recursion.
a is n x m and b is m x p (lists of rows); pass p when m = 0. The ring is
Z, or Z/mod Z when mod is set. Inputs with a side of at most n0 (n0 >= 1)
use the classical product; odd sides are padded by one zero row or column
at each level. Cost: O(n^(log2 7)) ring operations for n x n inputs;
recursion depth O(log max(n, m, p)).
"""
if n0 < 1:
raise ValueError("the cutoff n0 must be at least 1")
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")
return _strassen(a, b, n, m, p, mod, n0)
def _strassen(a: Matrix, b: Matrix, n: int, m: int, p: int, mod: int | None, n0: int) -> Matrix:
if min(n, m, p) <= n0: # @line L1
return _classical(a, b, p, mod)
if n % 2 or m % 2 or p % 2: # @line L2
n2, m2, p2 = n + n % 2, m + m % 2, p + p % 2
ap = [row + [0] * (m2 - m) for row in a] + [[0] * m2 for _ in range(n2 - n)]
bp = [row + [0] * (p2 - p) for row in b] + [[0] * p2 for _ in range(m2 - m)]
cp = _strassen(ap, bp, n2, m2, p2, mod, n0)
return [row[:p] for row in cp[:n]]
h, k, q = n // 2, m // 2, p // 2 # @line L3
a11, a12 = _block(a, 0, h, 0, k), _block(a, 0, h, k, m)
a21, a22 = _block(a, h, n, 0, k), _block(a, h, n, k, m)
b11, b12 = _block(b, 0, k, 0, q), _block(b, 0, k, q, p)
b21, b22 = _block(b, k, m, 0, q), _block(b, k, m, q, p)
def rec(x: Matrix, y: Matrix) -> Matrix:
return _strassen(x, y, h, k, q, mod, n0)
m1 = rec(_add(a11, a22, mod), _add(b11, b22, mod)) # @line L4
m2 = rec(_add(a21, a22, mod), b11) # @line L5
m3 = rec(a11, _sub(b12, b22, mod)) # @line L6
m4 = rec(a22, _sub(b21, b11, mod)) # @line L7
m5 = rec(_add(a11, a12, mod), b22) # @line L8
m6 = rec(_sub(a21, a11, mod), _add(b11, b12, mod)) # @line L9
m7 = rec(_sub(a12, a22, mod), _add(b21, b22, mod)) # @line L10
c11 = _add(_sub(_add(m1, m4, mod), m5, mod), m7, mod) # @line L11
c12 = _add(m3, m5, mod) # @line L12
c21 = _add(m2, m4, mod) # @line L13
c22 = _add(_add(_sub(m1, m2, mod), m3, mod), m6, mod) # @line L14
return [r + s for r, s in zip(c11, c12)] + [r + s for r, s in zip(c21, c22)] # @line L15
def count_operations(n: int, n0: int = 1) -> int:
"""Ring operations of the method on n x n inputs with n a power of two: 7 T(n/2) + 18 (n/2)^2."""
if n <= n0:
return 2 * n**3 - n**2
return 7 * count_operations(n // 2, n0) + 18 * (n // 2) ** 2
test_strassen_algorithm.py 86 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/strassen-algorithm/. The authors are listed in prov/MZ8PBV.jsonl.
import math
import random
import pytest
from strassen_algorithm import count_operations, strassen_algorithm 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:
assert mul([[1, 2], [3, 4]], [[5, 6], [7, 8]]) == [[19, 22], [43, 50]]
def test_edges() -> None:
assert mul([], []) == []
assert mul([[3]], [[4]]) == [[12]]
assert mul([[], []], [], p=2) == [[0, 0], [0, 0]]
with pytest.raises(ValueError):
mul([[1]], [[1]], n0=0)
with pytest.raises(ValueError):
mul([[1, 2]], [[1, 2]])
def test_identity_zero() -> None:
rng = random.Random(3)
for n in range(0, 12):
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(a, zero, p=n) == zero
def test_square_all_sizes_over_z() -> None:
rng = random.Random(1969)
for n in range(1, 21):
for _ in range(3):
a, b = rand(rng, n, n, -10**9, 10**9), rand(rng, n, n, -10**9, 10**9)
assert mul(a, b) == schoolbook(a, b, n)
def test_random_shapes_and_cutoffs() -> None:
rng = random.Random(4242)
for _ in range(500):
n, m, p = rng.randint(0, 10), rng.randint(0, 10), rng.randint(0, 10)
a, b = rand(rng, n, m), rand(rng, m, p)
assert mul(a, b, p=p, n0=rng.choice((1, 2, 3, 5))) == schoolbook(a, b, p)
def test_mod_p() -> None:
rng = random.Random(17)
for q in (2, 3, 5, 7, 11):
for _ in range(40):
n = rng.randint(1, 12)
a, b = rand(rng, n, n, 0, q - 1), rand(rng, n, n, 0, q - 1)
got = mul(a, b, mod=q)
assert got == schoolbook(a, b, n, q)
assert all(0 <= x < q for row in got for x in row)
def test_counts() -> None:
for k in range(0, 11):
n = 2**k
assert count_operations(n) == 7 * 7**k - 6 * n * n
assert round(7 ** k) == round(n ** math.log2(7))
assert count_operations(16, 8) == 7872 < 2 * 16**3 - 16**2
assert count_operations(1024) < 2 * 1024**3 - 1024**2 < count_operations(512) * 8
assert count_operations(512) > 2 * 512**3 - 512**2
What the next level needs
- the tests did not run in this build