Heapsort in Python
Builds a max-heap and repeatedly extracts the maximum; Theta(n log n) worst case and in place.
- File
heapsort.py, 41 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/heapsort/. The authors are listed in prov/CSPR1S.jsonl.
from typing import Protocol, TypeVar
class Ordered(Protocol):
def __lt__(self, other: "Ordered", /) -> bool: ...
T = TypeVar("T", bound=Ordered)
def heapsort(a: list[T]) -> None:
"""Sort a in place, in non-decreasing order.
The items need only the < operator. The sort is not stable.
It makes O(n log n) comparisons and uses O(1) extra space.
"""
n = len(a)
for i in range(n // 2 - 1, -1, -1): # @line L1
sift_down(a, i, n)
for end in range(n - 1, 0, -1): # @line L2
a[0], a[end] = a[end], a[0] # @line L3
sift_down(a, 0, end) # @line L4
def sift_down(a: list[T], i: int, m: int) -> None:
"""Move a[i] down in the max-heap a[0:m] until no child is larger.
Precondition: the subtrees of the children of i are max-heaps.
Cost: O(log m) comparisons and swaps.
"""
while 2 * i + 1 < m: # @line L5
c = 2 * i + 1 # @line L6
if c + 1 < m and a[c] < a[c + 1]: # @line L7
c += 1
if not a[i] < a[c]: # @line L8: a[i] >= a[c]
return
a[i], a[c] = a[c], a[i] # @line L9
i = c # @line L10
test_heapsort.py 57 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/heapsort/. The authors are listed in prov/CSPR1S.jsonl.
import itertools
import random
from heapsort import heapsort, sift_down
def test_example() -> None:
a = [6, 5, 3, 1, 8, 7, 2, 4]
heapsort(a)
assert a == [1, 2, 3, 4, 5, 6, 7, 8]
def test_build_heap_of_example() -> None:
a = [6, 5, 3, 1, 8, 7, 2, 4]
for i in range(len(a) // 2 - 1, -1, -1):
sift_down(a, i, len(a))
assert a == [8, 6, 7, 4, 5, 3, 2, 1]
def test_edges() -> None:
for a in ([], [1], [2, 1], [1, 2], [1, 1, 1], [3, 2, 1], [5, 5, 1, 5], [-(2**63), 2**63 - 1, 0]):
b = list(a)
heapsort(b)
assert b == sorted(a)
def test_all_permutations() -> None:
for n in range(7):
for p in itertools.permutations(range(n)):
b = list(p)
heapsort(b)
assert b == list(range(n))
def test_not_stable() -> None:
class Key:
def __init__(self, k: int, tag: str) -> None:
self.k, self.tag = k, tag
def __lt__(self, other: "Key") -> bool:
return self.k < other.k
a = [Key(1, "a"), Key(1, "b")]
heapsort(a)
assert [x.tag for x in a] == ["b", "a"]
def test_random_against_sorted() -> None:
rng = random.Random(1964)
for _ in range(600):
a = [rng.randint(-40, 40) for _ in range(rng.randint(0, 70))]
b = list(a)
heapsort(b)
assert b == sorted(a)
What the next level needs
- the tests did not run in this build