Median of medians in Python
Deterministic linear-time selection that uses the median of group medians as the pivot.
- File
median_of_medians.py, 55 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/median-of-medians/. The authors are listed in prov/A5DV6Z.jsonl.
from typing import Protocol, TypeVar
class Ordered(Protocol):
def __lt__(self, other: "Ordered", /) -> bool: ...
T = TypeVar("T", bound=Ordered)
def median_of_medians(s: list[T], k: int) -> T:
"""Return the k-th smallest item of s, for 1 <= k <= len(s).
Deterministic selection of Blum, Floyd, Pratt, Rivest and Tarjan with groups
of 5. The input list is not changed. Time Theta(n) in the worst case, extra
space Theta(n) for the lists of the recursive calls; the recursion depth is
O(log n). Only the operator < is used on the items.
"""
n = len(s)
if not 1 <= k <= n:
raise ValueError("k must satisfy 1 <= k <= len(s)")
if n <= 5: # @line L1
return insertion_sorted(s)[k - 1]
groups = [s[i : i + 5] for i in range(0, n, 5)] # @line L2
medians = [lower_median(g) for g in groups] # @line L3
x = median_of_medians(medians, (len(medians) + 1) // 2) # @line L4
less = [v for v in s if v < x] # @line L5
greater = [v for v in s if x < v]
equal = n - len(less) - len(greater)
if k <= len(less): # @line L6
return median_of_medians(less, k)
if k <= len(less) + equal: # @line L7
return x
return median_of_medians(greater, k - len(less) - equal) # @line L8
def insertion_sorted(g: list[T]) -> list[T]:
"""Return a sorted copy of a short list, by insertion sort."""
b = list(g)
for i in range(1, len(b)):
v = b[i]
j = i - 1
while j >= 0 and v < b[j]:
b[j + 1] = b[j]
j -= 1
b[j + 1] = v
return b
def lower_median(g: list[T]) -> T:
"""Return the item of rank ceil(len(g) / 2) of a group of at most 5 items."""
return insertion_sorted(g)[(len(g) + 1) // 2 - 1]
test_median_of_medians.py 59 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/median-of-medians/. The authors are listed in prov/A5DV6Z.jsonl.
import random
import pytest
from median_of_medians import lower_median, median_of_medians
def test_example() -> None:
a = [12, 3, 45, 7, 29, 18, 1, 33, 26, 9, 40, 15, 21, 4, 37]
assert median_of_medians(a, 8) == 18
assert [median_of_medians(a, k) for k in range(1, 16)] == sorted(a)
def test_group_medians() -> None:
assert lower_median([12, 3, 45, 7, 29]) == 12
assert lower_median([4, 1]) == 1
assert lower_median([9]) == 9
def test_edges() -> None:
assert median_of_medians([5], 1) == 5
assert median_of_medians([2, 2, 2, 2, 2, 2, 2], 4) == 2
assert median_of_medians(list(range(100, 0, -1)), 1) == 1
assert median_of_medians(list(range(100)), 100) == 99
with pytest.raises(ValueError):
median_of_medians([], 1)
with pytest.raises(ValueError):
median_of_medians([1], 0)
def test_input_unchanged() -> None:
a = [3, 1, 2, 5, 4, 9, 7, 8, 6]
b = list(a)
median_of_medians(a, 5)
assert a == b
def test_pivot_split_bound() -> None:
rng = random.Random(55)
for _ in range(300):
n = rng.randint(6, 200)
s = [rng.randint(0, 30) for _ in range(n)]
meds = [lower_median(s[i : i + 5]) for i in range(0, n, 5)]
x = sorted(meds)[(len(meds) + 1) // 2 - 1]
less = sum(1 for v in s if v < x)
greater = sum(1 for v in s if v > x)
assert 10 * less <= 7 * n + 20 and 10 * greater <= 7 * n + 20
def test_random_against_sorted() -> None:
rng = random.Random(4417)
for _ in range(600):
n = rng.randint(1, 120)
a = [rng.randint(-40, 40) for _ in range(n)]
k = rng.randint(1, n)
assert median_of_medians(a, k) == sorted(a)[k - 1]
What the next level needs
- the tests did not run in this build