Monotonic deque over prefix sums in Python
Keep an increasing deque of prefix sums to find, for each right end, the nearest valid left end in amortized O(1).
- File
monotonic_deque_prefix_sums.py, 64 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/monotonic-deque-prefix-sums/. The authors are listed in prov/HHFZR5.jsonl.
from collections import deque
from typing import Optional, Sequence, Union
Number = Union[int, float]
def shortest_subarray_at_least(a: Sequence[Number], k: Number) -> Optional[tuple[int, int]]:
"""Return (i, j) such that a[i:j] is a shortest non-empty subarray with sum >= k, or None.
Items may be negative. Among shortest pairs, the one with the smallest end is returned.
Cost: Theta(n) time and Theta(n) extra space.
"""
n = len(a)
p: list[Number] = [0] * (n + 1) # @line L1
for t in range(1, n + 1):
p[t] = p[t - 1] + a[t - 1]
q: deque[int] = deque() # @line L2
best = n + 1
r: Optional[tuple[int, int]] = None
for j in range(n + 1): # @line L3
while q and p[j] - p[q[0]] >= k: # @line L4
i = q.popleft() # @line L5
if j - i < best:
best, r = j - i, (i, j)
while q and p[q[-1]] >= p[j]: # @line L6
q.pop()
q.append(j) # @line L7
return r # @line L8
def shortest_length_at_least(a: Sequence[Number], k: Number) -> int:
"""Return the length of a shortest subarray with sum >= k, or -1 if there is none."""
r = shortest_subarray_at_least(a, k)
return -1 if r is None else r[1] - r[0]
def max_sum_with_length_between(a: Sequence[Number], lo: int, hi: int) -> Optional[Number]:
"""Return the largest sum of a[i:j] with lo <= j - i <= hi, or None if no length fits.
Precondition: 1 <= lo <= hi (raises ValueError otherwise).
Cost: Theta(n) time, a sliding window minimum of the prefix sums.
"""
if lo < 1 or hi < lo:
raise ValueError("need 1 <= lo <= hi")
n = len(a)
p: list[Number] = [0] * (n + 1)
for t in range(1, n + 1):
p[t] = p[t - 1] + a[t - 1]
q: deque[int] = deque()
best: Optional[Number] = None
for j in range(lo, n + 1):
s = j - lo
while q and p[q[-1]] >= p[s]:
q.pop()
q.append(s)
while q[0] < j - hi:
q.popleft()
cand = p[j] - p[q[0]]
if best is None or cand > best:
best = cand
return best
test_monotonic_deque_prefix_sums.py 66 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/monotonic-deque-prefix-sums/. The authors are listed in prov/HHFZR5.jsonl.
import random
import pytest
from monotonic_deque_prefix_sums import (
max_sum_with_length_between,
shortest_length_at_least,
shortest_subarray_at_least,
)
def brute_shortest(a: list[int], k: int) -> int:
n = len(a)
best = -1
for i in range(n):
s = 0
for j in range(i, n):
s += a[j]
if s >= k and (best < 0 or j - i + 1 < best):
best = j - i + 1
return best
def brute_band(a: list[int], lo: int, hi: int):
sums = [sum(a[i:j]) for i in range(len(a)) for j in range(i + lo, min(len(a), i + hi) + 1)]
return max(sums) if sums else None
def test_example() -> None:
assert shortest_subarray_at_least([84, -37, 32, 40, 95], 167) == (2, 5)
assert shortest_length_at_least([2, -3, 4, -1, 3, -5, 6], 6) == 1
assert shortest_length_at_least([2, -1, 2], 3) == 3
assert shortest_length_at_least([1, 2], 4) == -1
def test_edges() -> None:
assert shortest_subarray_at_least([], 0) is None
assert shortest_subarray_at_least([-5], 0) is None
assert shortest_subarray_at_least([-5], -5) == (0, 1)
assert shortest_length_at_least([1, 1, 1], -10) == 1
assert max_sum_with_length_between([1, -2, 3], 2, 3) == 2
assert max_sum_with_length_between([1, 2], 3, 4) is None
with pytest.raises(ValueError):
max_sum_with_length_between([1], 0, 1)
def test_against_brute_force() -> None:
rng = random.Random(20261105)
for _ in range(3000):
n = rng.randrange(0, 12)
a = [rng.randrange(-8, 9) for _ in range(n)]
k = rng.randrange(-6, 20)
expected = brute_shortest(a, k)
r = shortest_subarray_at_least(a, k)
if expected < 0:
assert r is None
else:
assert r is not None
i, j = r
assert j - i == expected and sum(a[i:j]) >= k
lo = rng.randrange(1, 5)
hi = lo + rng.randrange(0, 5)
assert max_sum_with_length_between(a, lo, hi) == brute_band(a, lo, hi)
What the next level needs
- the tests did not run in this build