Stack with minimum in Python
A stack that also returns its current minimum element in constant time by storing the running minimum with each entry.
- File
min_stack.py, 85 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/min-stack/. The authors are listed in prov/PR4H8X.jsonl.
from typing import Generic, Protocol, TypeVar
class Ordered(Protocol):
def __lt__(self, other: "Ordered", /) -> bool: ...
T = TypeVar("T", bound=Ordered)
class MinStack(Generic[T]):
"""A stack that also returns the least item it holds.
Each entry is a pair (x, m), where m is the least item at or below that
entry. push, pop, top and min cost O(1) amortized on a Python list
(O(1) worst case on a fixed array or a linked list). Space: 2n cells.
"""
def __init__(self) -> None:
self.s: list[tuple[T, T]] = []
def __len__(self) -> int:
return len(self.s)
def is_empty(self) -> bool:
return not self.s
def push(self, x: T) -> None:
"""Push x. Cost O(1)."""
if not self.s: # @line L1
m = x
else:
m = self.s[-1][1]
if x < m:
m = x
self.s.append((x, m)) # @line L2
def pop(self) -> T:
"""Remove and return the top item. Precondition: not empty. Cost O(1)."""
if not self.s:
raise IndexError("pop from an empty stack")
x, _ = self.s.pop() # @line L3
return x # @line L4
def top(self) -> T:
"""Return the top item. Precondition: not empty. Cost O(1)."""
if not self.s:
raise IndexError("top of an empty stack")
return self.s[-1][0] # @line L5
def min(self) -> T:
"""Return the least item. Precondition: not empty. Cost O(1)."""
if not self.s:
raise IndexError("min of an empty stack")
return self.s[-1][1] # @line L6
def check_invariant(self) -> None:
for k, (x, m) in enumerate(self.s):
below = self.s[k - 1][1] if k else x
assert not x < m and not below < m and (m == x or m == below)
class LeanMinStack(Generic[T]):
"""The variant with a second stack that holds only the new minima (pushed when x <= min)."""
def __init__(self) -> None:
self.items: list[T] = []
self.mins: list[T] = []
def push(self, x: T) -> None:
self.items.append(x)
if not self.mins or not self.mins[-1] < x: # x <= current min
self.mins.append(x)
def pop(self) -> T:
x = self.items.pop()
if not self.mins[-1] < x and not x < self.mins[-1]: # x == current min
self.mins.pop()
return x
def min(self) -> T:
return self.mins[-1]
test_min_stack.py 68 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/min-stack/. The authors are listed in prov/PR4H8X.jsonl.
import random
import pytest
from min_stack import LeanMinStack, MinStack
def test_page_example() -> None:
s: MinStack[int] = MinStack()
mins = []
for x in [5, 3, 7, 3, 2]:
s.push(x)
mins.append(s.min())
assert mins == [5, 3, 3, 3, 2]
assert s.s == [(5, 5), (3, 3), (7, 3), (3, 3), (2, 2)]
out = []
for _ in range(4):
out.append((s.pop(), s.min()))
assert out == [(2, 3), (3, 3), (7, 3), (3, 5)]
assert s.top() == 5
def test_edges() -> None:
s: MinStack[int] = MinStack()
for f in (s.pop, s.top, s.min):
with pytest.raises(IndexError):
f()
s.push(4)
s.push(4)
s.pop()
assert s.min() == 4 # a duplicate of the minimum
s.push(-(2**63))
s.push(2**63 - 1)
assert s.min() == -(2**63)
def test_lean_variant_with_duplicates() -> None:
s: LeanMinStack[int] = LeanMinStack()
for x in [2, 1, 1]:
s.push(x)
s.pop()
assert s.min() == 1 # with < instead of <= this would be 2
def test_random_against_model() -> None:
rng = random.Random(155)
for _ in range(500):
s: MinStack[int] = MinStack()
t: LeanMinStack[int] = LeanMinStack()
model: list[int] = []
for _ in range(50):
if model and rng.random() < 0.45:
want = model.pop()
assert s.pop() == want
assert t.pop() == want
else:
x = rng.randint(0, 6)
s.push(x)
t.push(x)
model.append(x)
s.check_invariant()
if model:
assert s.min() == min(model) == t.min()
assert s.top() == model[-1]
assert len(s) == len(model)
What the next level needs
- the tests did not run in this build