Sink (sift-down) operation in Python
Restores heap order by exchanging a node with its larger child repeatedly until the node is no smaller than its children.
- File
sift_down_sink.py, 72 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/sift-down-sink/. The authors are listed in prov/RCFB8N.jsonl.
from collections.abc import Callable
from typing import Any, Protocol, TypeVar
class Ordered(Protocol):
def __lt__(self, other: Any, /) -> bool: ...
T = TypeVar("T", bound=Ordered)
def sink(less: Callable[[int, int], bool], exch: Callable[[int, int], None], k: int, n: int) -> None:
"""Sink the item at position k of a 1-based max-heap of n items.
less(i, j) says whether the item at position i is smaller than the item at
position j, and exch(i, j) swaps them. Precondition: the subtrees of the
children 2k and 2k + 1 are max-heaps. Postcondition: the subtree of k is a
max-heap. Cost: at most 2 * floor(log2(n / k)) calls of less.
"""
while 2 * k <= n: # @line L1
j = 2 * k # @line L2
if j < n and less(j, j + 1): # @line L3
j += 1
if not less(k, j): # @line L4
break
exch(k, j) # @line L5
k = j # @line L6
def del_max(pq: list[T]) -> T:
"""Remove and return a largest item of the max-heap pq[1:]; pq[0] is not used. O(log n)."""
n = len(pq) - 1
if n < 1:
raise IndexError("del_max from an empty priority queue")
def less(i: int, j: int) -> bool:
return pq[i] < pq[j]
def exch(i: int, j: int) -> None:
pq[i], pq[j] = pq[j], pq[i]
top = pq[1] # @line L7
exch(1, n) # @line L8
n -= 1
sink(less, exch, 1, n) # @line L9
pq.pop() # @line L10: drop the old maximum from the last cell
return top # @line L11
def heapsort(a: list[T]) -> None:
"""Sort a in place in non-decreasing order with heap construction by sink and the sortdown.
The heap uses positions 1 to n, and position p is the cell a[p - 1].
Cost: fewer than 2n lg n + 2n comparisons, O(1) extra space.
"""
def less(i: int, j: int) -> bool:
return a[i - 1] < a[j - 1]
def exch(i: int, j: int) -> None:
a[i - 1], a[j - 1] = a[j - 1], a[i - 1]
n = len(a)
for k in range(n // 2, 0, -1): # @line L12
sink(less, exch, k, n)
while n > 1: # @line L13
exch(1, n) # @line L14
n -= 1
sink(less, exch, 1, n) # @line L15
test_sift_down_sink.py 85 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/sift-down-sink/. The authors are listed in prov/RCFB8N.jsonl.
import random
import pytest
from sift_down_sink import del_max, heapsort, sink
def is_max_heap(pq: list[int]) -> bool:
n = len(pq) - 1
return all(not pq[k // 2] < pq[k] for k in range(2, n + 1))
def counting(a: list[int]) -> tuple[list[int], object, object]:
calls = [0]
def less(i: int, j: int) -> bool:
calls[0] += 1
return a[i - 1] < a[j - 1]
def exch(i: int, j: int) -> None:
a[i - 1], a[j - 1] = a[j - 1], a[i - 1]
return calls, less, exch
def test_example_construction() -> None:
a = [5, 1, 9, 3, 7, 6, 2, 8, 4]
calls, less, exch = counting(a)
for k in range(len(a) // 2, 0, -1):
sink(less, exch, k, len(a)) # type: ignore[arg-type]
assert a == [9, 8, 6, 4, 7, 5, 2, 3, 1]
assert calls[0] == 12
def test_example_heapsort() -> None:
a = [5, 1, 9, 3, 7, 6, 2, 8, 4]
heapsort(a)
assert a == [1, 2, 3, 4, 5, 6, 7, 8, 9]
def test_del_max() -> None:
pq = [None, 9, 8, 6, 4, 7, 5, 2, 3, 1]
assert del_max(pq) == 9 # type: ignore[arg-type]
assert pq[1:] == [8, 7, 6, 4, 1, 5, 2, 3]
assert is_max_heap(pq) # type: ignore[arg-type]
def test_edges() -> None:
with pytest.raises(IndexError):
del_max([None]) # type: ignore[list-item]
pq = [None, 4]
assert del_max(pq) == 4 and pq == [None] # type: ignore[arg-type]
for a in ([], [1], [2, 1], [1, 1, 1], [3, 2, 1], [1, 2, 3]):
b = list(a)
heapsort(b)
assert b == sorted(a)
def test_random_against_sorted() -> None:
rng = random.Random(2024)
for _ in range(600):
a = [rng.randint(-40, 40) for _ in range(rng.randint(0, 50))]
b = list(a)
heapsort(b)
assert b == sorted(a)
def test_random_del_max() -> None:
rng = random.Random(5)
for _ in range(300):
items = [rng.randint(0, 30) for _ in range(rng.randint(1, 40))]
a = list(items)
calls, less, exch = counting(a)
for k in range(len(a) // 2, 0, -1):
sink(less, exch, k, len(a)) # type: ignore[arg-type]
assert calls[0] < 2 * len(a)
pq = [None] + a
out = []
while len(pq) > 1:
out.append(del_max(pq)) # type: ignore[arg-type]
assert is_max_heap(pq) # type: ignore[arg-type]
assert out == sorted(items, reverse=True)
What the next level needs
- the tests did not run in this build