SMAWK algorithm in Python
Find all row minima of a totally monotone n by m matrix in O(n + m) evaluations.
- File
smawk_algorithm.py, 50 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/smawk-algorithm/. The authors are listed in prov/V8AHMD.jsonl.
from typing import Callable, Sequence
def smawk(rows: Sequence[int], cols: Sequence[int], a: Callable[[int, int], float]) -> dict[int, int]:
"""Return a map from each row to the column of its leftmost minimum.
rows and cols are increasing lists of indices, cols is not empty, and
a(i, j) returns the entry in row i and column j. The matrix restricted to
rows and cols must be totally monotone: for i < i2 and j < j2,
a(i, j) > a(i, j2) implies a(i2, j) > a(i2, j2).
It reads O(len(rows) + len(cols)) entries and takes the same time.
"""
if not rows: # @line L1
return {}
c = reduce_columns(rows, cols, a) # @line L2
odd = rows[1::2] # @line L3
m = smawk(odd, c, a) # @line L4
pos = {col: t for t, col in enumerate(c)}
k = 0 # @line L5
for i in range(0, len(rows), 2): # @line L6
stop = pos[m[rows[i + 1]]] if i + 1 < len(rows) else len(c) - 1 # @line L7
r = rows[i]
best, best_val = c[k], a(r, c[k]) # @line L8
for t in range(k + 1, stop + 1):
v = a(r, c[t])
if v < best_val:
best, best_val = c[t], v
m[r] = best
k = stop # @line L9
return m # @line L10
def reduce_columns(rows: Sequence[int], cols: Sequence[int], a: Callable[[int, int], float]) -> list[int]:
"""Drop columns that hold no leftmost row minimum; at most len(rows) columns stay."""
s: list[int] = [] # @line L11
for c in cols: # @line L12
while s and a(rows[len(s) - 1], s[-1]) > a(rows[len(s) - 1], c): # @line L13
s.pop()
if len(s) < len(rows): # @line L14
s.append(c)
return s # @line L15
def row_minima(n: int, m: int, a: Callable[[int, int], float]) -> list[int]:
"""Leftmost minimum column of each row of a totally monotone n by m matrix (m >= 1)."""
res = smawk(list(range(n)), list(range(m)), a)
return [res[i] for i in range(n)]
test_smawk_algorithm.py 98 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/smawk-algorithm/. The authors are listed in prov/V8AHMD.jsonl.
import random
from smawk_algorithm import row_minima
def leftmost_minima(mat):
return [row.index(min(row)) for row in mat]
def is_totally_monotone(mat):
n, m = len(mat), len(mat[0])
for i in range(n):
for i2 in range(i + 1, n):
for j in range(m):
for j2 in range(j + 1, m):
if mat[i][j] > mat[i][j2] and not mat[i2][j] > mat[i2][j2]:
return False
return True
def run(mat):
count = [0]
def a(i, j):
count[0] += 1
return mat[i][j]
return row_minima(len(mat), len(mat[0]), a), count[0]
def random_monge(rng, n, m):
# Sums of Monge terms are Monge: column and row offsets, and f(x_i - y_j) with f convex.
xs = sorted(rng.randint(0, 20) for _ in range(n))
ys = sorted(rng.randint(0, 20) for _ in range(m))
u = [rng.randint(0, 9) for _ in range(m)]
v = [rng.randint(0, 9) for _ in range(n)]
return [[(xs[i] - ys[j]) ** 2 + u[j] + v[i] for j in range(m)] for i in range(n)]
def test_example() -> None:
x, y, u = [1, 2, 4, 5, 7], range(7), [9, 4, 6, 0, 7, 3, 5]
mat = [[(xi - yj) ** 2 + uj for yj, uj in zip(y, u)] for xi in x]
assert run(mat)[0] == [1, 3, 3, 5, 6]
def test_edges() -> None:
assert run([[5]])[0] == [0]
assert run([[3, 1, 2]])[0] == [1]
assert run([[2], [1], [7]])[0] == [0, 0, 0]
assert run([[4, 4, 4], [4, 4, 4]])[0] == [0, 0]
assert row_minima(0, 3, lambda i, j: 0) == []
def test_ties_take_leftmost() -> None:
mat = [[1, 1, 2], [3, 0, 0], [5, 2, 2]]
assert is_totally_monotone(mat)
assert run(mat)[0] == leftmost_minima(mat)
def test_large_values() -> None:
big = 1 << 62
mat = [[big, big - 1, big], [big + 1, big, big - 5]]
assert is_totally_monotone(mat)
assert run(mat)[0] == [1, 2]
def test_random_monge_against_brute_force() -> None:
rng = random.Random(1987)
for _ in range(800):
n, m = rng.randint(1, 12), rng.randint(1, 12)
mat = random_monge(rng, n, m)
got, reads = run(mat)
assert got == leftmost_minima(mat)
assert reads <= 4 * m + 12 * n
def test_random_totally_monotone_small_values() -> None:
rng = random.Random(42)
done = 0
while done < 600:
n, m = rng.randint(1, 4), rng.randint(1, 4)
mat = [[rng.randint(0, 3) for _ in range(m)] for _ in range(n)]
if not is_totally_monotone(mat):
continue
done += 1
assert run(mat)[0] == leftmost_minima(mat)
def test_linear_reads_on_large_input() -> None:
rng = random.Random(7)
n, m = 300, 2000
mat = random_monge(rng, n, m)
got, reads = run(mat)
assert got == leftmost_minima(mat)
assert reads <= 4 * m + 12 * n
What the next level needs
- the tests did not run in this build