Bentley's maximum sum submatrix algorithm in Python

Fix a pair of rows, compress columns into a 1D array, then run Kadane; runs in O(n^3) for an n by n grid.

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

bentley_column_sweep.pyRaw GitHub

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/bentley-column-sweep/. The authors are listed in prov/B4PWT1.jsonl.

from typing import Optional, Union

Number = Union[int, float]


def _kadane(c: list[Number]) -> tuple[Number, int, int]:
    """Return (best, e, d): the largest sum of a non-empty subarray of c and its columns."""
    cur = c[0]
    start = 0
    best = c[0]
    be, bd = 0, 0
    for j in range(1, len(c)):
        if cur < 0:
            cur, start = c[j], j
        else:
            cur += c[j]
        if cur > best:
            best, be, bd = cur, start, j
    return best, be, bd


def max_sum_submatrix(m: list[list[Number]]) -> tuple[Number, tuple[int, int, int, int]]:
    """Return (best, (a, b, e, d)): rows a..b and columns e..d of a largest-sum rectangle.

    Indices are 0-based and inclusive. Precondition: m is a non-empty rectangular
    matrix with at least one column (raises ValueError otherwise).
    Cost: Theta(r^2 c) time for r rows and c columns, Theta(c) extra space. The rows
    are swept as given; transpose first if there are more rows than columns.
    """
    if not m or not m[0] or any(len(row) != len(m[0]) for row in m):
        raise ValueError("the matrix must be non-empty and rectangular")
    rows, cols = len(m), len(m[0])
    best: Optional[Number] = None  # @line L1
    rect = (0, 0, 0, 0)
    for a in range(rows):  # @line L2
        c: list[Number] = [0] * cols  # @line L3
        for b in range(a, rows):  # @line L4
            row = m[b]  # @line L5
            for j in range(cols):
                c[j] += row[j]
            s, e, d = _kadane(c)  # @line L6
            if best is None or s > best:  # @line L7
                best, rect = s, (a, b, e, d)
    assert best is not None
    return best, rect  # @line L8


def max_sum_submatrix_auto(m: list[list[Number]]) -> Number:
    """Return the largest rectangle sum, sweeping bands of the shorter side.

    Cost: Theta(min(r, c)^2 max(r, c)) time.
    """
    if not m or not m[0]:
        raise ValueError("the matrix must be non-empty and rectangular")
    if len(m) > len(m[0]):
        m = [list(col) for col in zip(*m)]
    return max_sum_submatrix(m)[0]

test_bentley_column_sweep.py 63 lines

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/bentley-column-sweep/. The authors are listed in prov/B4PWT1.jsonl.

import random

import pytest

from bentley_column_sweep import max_sum_submatrix, max_sum_submatrix_auto


def rect_sum(m: list[list[int]], a: int, b: int, e: int, d: int) -> int:
    return sum(m[i][j] for i in range(a, b + 1) for j in range(e, d + 1))


def brute(m: list[list[int]]) -> int:
    r, c = len(m), len(m[0])
    return max(
        rect_sum(m, a, b, e, d)
        for a in range(r)
        for b in range(a, r)
        for e in range(c)
        for d in range(e, c)
    )


EXAMPLE = [
    [1, 2, -1, -4, -20],
    [-8, -3, 4, 2, 1],
    [3, 8, 10, 1, 3],
    [-4, -1, 1, 7, -6],
]


def test_example() -> None:
    assert max_sum_submatrix(EXAMPLE) == (29, (1, 3, 1, 3))
    assert max_sum_submatrix_auto(EXAMPLE) == 29


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


def test_against_brute_force() -> None:
    rng = random.Random(20261105)
    for _ in range(2000):
        r = rng.randrange(1, 6)
        c = rng.randrange(1, 6)
        m = [[rng.randrange(-9, 10) for _ in range(c)] for _ in range(r)]
        expected = brute(m)
        best, (a, b, e, d) = max_sum_submatrix(m)
        assert best == expected
        assert 0 <= a <= b < r and 0 <= e <= d < c
        assert rect_sum(m, a, b, e, d) == best
        assert max_sum_submatrix_auto(m) == expected

What the next level needs

the tests did not run in this build

Read the full page about Bentley's maximum sum submatrix algorithm