Prefix sum hash counting in Python
Store prefix-sum frequencies in a hash table; each subarray sum equals a difference of two prefix sums.
- File
prefix_sum_hash_counting.py, 54 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/prefix-sum-hash-counting/. The authors are listed in prov/ASHEE3.jsonl.
def count_subarrays_with_sum(a: list[int], k: int) -> int:
"""Return the number of non-empty subarrays of a whose sum is exactly k.
Items may be negative. Precondition: integer items, so that sums compare exactly.
Cost: O(n) expected time with Python's dict, O(n) extra space.
"""
counts: dict[int, int] = {0: 1} # @line L1
p = 0
count = 0
for x in a: # @line L2
p += x # @line L3
count += counts.get(p - k, 0) # @line L4
counts[p] = counts.get(p, 0) + 1 # @line L5
return count # @line L6
def longest_subarray_with_sum(a: list[int], k: int) -> tuple[int, int]:
"""Return (length, start) of a longest subarray with sum k, or (0, -1) if there is none.
Among the longest ones, the one with the earliest end is returned.
Cost: O(n) expected time, O(n) extra space.
"""
first: dict[int, int] = {0: 0}
p = 0
best, start = 0, -1
for j, x in enumerate(a):
p += x
i = first.get(p - k)
if i is not None and j + 1 - i > best:
best, start = j + 1 - i, i
if p not in first:
first[p] = j + 1
return best, start
def count_subarrays_divisible(a: list[int], k: int) -> int:
"""Return the number of non-empty subarrays of a whose sum is a multiple of k > 0.
Cost: O(n + k) time, O(k) extra space.
"""
if k <= 0:
raise ValueError("k must be positive")
counts = [0] * k
counts[0] = 1
p = 0
total = 0
for x in a:
p = (p + x) % k
total += counts[p]
counts[p] += 1
return total
test_prefix_sum_hash_counting.py 51 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/prefix-sum-hash-counting/. The authors are listed in prov/ASHEE3.jsonl.
import random
import pytest
from prefix_sum_hash_counting import (
count_subarrays_divisible,
count_subarrays_with_sum,
longest_subarray_with_sum,
)
def sums(a: list[int]) -> list[tuple[int, int, int]]:
return [(i, j, sum(a[i : j + 1])) for i in range(len(a)) for j in range(i, len(a))]
def test_example() -> None:
assert count_subarrays_with_sum([3, 4, 7, 2, -3, 1, 4, 2], 7) == 4
assert count_subarrays_with_sum([1, 1, 1], 2) == 2
assert longest_subarray_with_sum([15, -2, 2, -8, 1, 7, 10, 23], 0) == (5, 1)
def test_edges() -> None:
assert count_subarrays_with_sum([], 0) == 0
assert count_subarrays_with_sum([7], 7) == 1
assert count_subarrays_with_sum([0, 0, 0, 0], 0) == 10
assert count_subarrays_with_sum([1, 2, 3], 7) == 0
assert longest_subarray_with_sum([1, 2, 3], 7) == (0, -1)
assert count_subarrays_divisible([4, 5, 0, -2, -3, 1], 5) == 7
assert count_subarrays_divisible([-3], 3) == 1
with pytest.raises(ValueError):
count_subarrays_divisible([1], 0)
def test_against_brute_force() -> None:
rng = random.Random(20261105)
for _ in range(2000):
n = rng.randrange(0, 20)
a = [rng.randrange(-5, 6) for _ in range(n)]
k = rng.randrange(-6, 7)
all_sums = sums(a)
assert count_subarrays_with_sum(a, k) == sum(1 for _, _, s in all_sums if s == k)
lengths = [j - i + 1 for i, j, s in all_sums if s == k]
length, start = longest_subarray_with_sum(a, k)
assert length == max(lengths, default=0)
if length:
assert sum(a[start : start + length]) == k
m = rng.randrange(1, 7)
assert count_subarrays_divisible(a, m) == sum(1 for _, _, s in all_sums if s % m == 0)
What the next level needs
- the tests did not run in this build