Schroeppel-Shamir algorithm in Python
Subset sum in O(2^(n/2)) time with only O(2^(n/4)) memory, using heaps to generate the sums in order.
- File
schroeppel_shamir_algorithm.py, 49 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/schroeppel-shamir-algorithm/. The authors are listed in prov/9AEYS5.jsonl.
import heapq
def _subset_sums(a: list[int], lo: int, hi: int) -> list[tuple[int, int]]:
"""Return (sum, mask) for each of the 2^(hi - lo) subsets of positions lo .. hi - 1."""
sums = [(0, 0)]
for p in range(lo, hi):
sums += [(s + a[p], m | (1 << p)) for s, m in sums]
return sums
def schroeppel_shamir(a: list[int], t: int) -> list[int] | None:
"""Return the sorted positions of a subset of a with sum t, or None if there is none.
The items are integers of any sign. The four quarter lists and the two
heaps hold O(2^(n/4)) entries; the walk does O(2^(n/2)) heap operations of
O(n) comparisons each, so the time is O(n 2^(n/2)).
"""
n = len(a)
q = [k * n // 4 for k in range(5)]
la = _subset_sums(a, q[0], q[1]) # @line L1
lb = _subset_sums(a, q[1], q[2])
lc = _subset_sums(a, q[2], q[3])
ld = _subset_sums(a, q[3], q[4])
lb.sort() # @line L2
ld.sort()
h1 = [(la[i][0] + lb[0][0], i, 0) for i in range(len(la))] # @line L3
heapq.heapify(h1)
last = len(ld) - 1
h2 = [(-(lc[k][0] + ld[last][0]), k, last) for k in range(len(lc))] # @line L4
heapq.heapify(h2)
while h1 and h2: # @line L5
s1, i, j = h1[0] # @line L6
s2, k, l = -h2[0][0], h2[0][1], h2[0][2]
if s1 + s2 == t: # @line L7
mask = la[i][1] | lb[j][1] | lc[k][1] | ld[l][1]
return [p for p in range(n) if mask >> p & 1]
if s1 + s2 < t: # @line L8
heapq.heappop(h1)
if j + 1 < len(lb):
heapq.heappush(h1, (la[i][0] + lb[j + 1][0], i, j + 1))
else: # @line L9
heapq.heappop(h2)
if l > 0:
heapq.heappush(h2, (-(lc[k][0] + ld[l - 1][0]), k, l - 1))
return None # @line L10
test_schroeppel_shamir_algorithm.py 50 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/schroeppel-shamir-algorithm/. The authors are listed in prov/9AEYS5.jsonl.
import itertools
import random
from schroeppel_shamir_algorithm import schroeppel_shamir
def brute_force(a: list[int], t: int) -> bool:
return any(
sum(c) == t
for r in range(len(a) + 1)
for c in itertools.combinations(a, r)
)
def is_valid(a: list[int], t: int, s: list[int] | None) -> bool:
return s is not None and s == sorted(set(s)) and all(0 <= p < len(a) for p in s) and sum(a[p] for p in s) == t
def test_example() -> None:
a = [31, 17, 45, 8, 26, 53, 12, 39]
assert schroeppel_shamir(a, 110) == [0, 4, 5]
assert schroeppel_shamir(a, 2) is None
def test_edges() -> None:
assert schroeppel_shamir([], 0) == []
assert schroeppel_shamir([], 5) is None
assert schroeppel_shamir([7], 7) == [0]
assert schroeppel_shamir([7], 0) == []
assert schroeppel_shamir([7], 3) is None
assert is_valid([5, 5, 5, 5, 5], 15, schroeppel_shamir([5, 5, 5, 5, 5], 15))
assert schroeppel_shamir([5, 5, 5, 5, 5], 7) is None
big = [2**62, 2**62 - 1, 3, -(2**62), 2**61, 1]
assert is_valid(big, 2**62 + 2**61 + 2, schroeppel_shamir(big, 2**62 + 2**61 + 2))
def test_against_brute_force() -> None:
rng = random.Random(20240611)
for _ in range(800):
n = rng.randrange(0, 11)
a = [rng.randrange(-30, 60) for _ in range(n)]
t = rng.randrange(-40, 150)
got = schroeppel_shamir(a, t)
if brute_force(a, t):
assert is_valid(a, t, got)
else:
assert got is None
What the next level needs
- the tests did not run in this build