Heap extract-max in Python

A priority queue operation that removes and returns the largest element of a max-heap, then restores the heap property.

File: heap_extract_max.py, 40 lines
Level: I0 not checked yet
Toolchain: Python 3.14, pytest 9
Test framework: pytest
Run the tests: python3 -m pytest

heap_extract_max.pyRaw GitHub

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/heap-extract-max/. The authors are listed in prov/4GYGY9.jsonl.

from typing import Any, Protocol, TypeVar


class Ordered(Protocol):
    def __lt__(self, other: Any, /) -> bool: ...


T = TypeVar("T", bound=Ordered)


def max_heapify(a: list[T], n: int, i: int) -> None:
    """Sift a[i] down in a[0:n]. All edges of a[0:n] must hold except the ones out of i."""
    while True:  # @line L7
        left, right, m = 2 * i + 1, 2 * i + 2, i  # @line L8
        if left < n and a[m] < a[left]:  # @line L9
            m = left
        if right < n and a[m] < a[right]:  # @line L10
            m = right
        if m == i:  # @line L11
            return
        a[i], a[m] = a[m], a[i]  # @line L12
        i = m


def heap_extract_max(a: list[T]) -> T:
    """Remove and return the largest key of the max-heap a (0-based). O(log n) time.

    a must satisfy a[(j - 1) // 2] >= a[j] for each j >= 1; it stays a max-heap.
    """
    if len(a) < 1:  # @line L1
        raise IndexError("heap underflow")
    top = a[0]  # @line L2
    last = a.pop()  # @line L4
    if a:
        a[0] = last  # @line L3
        max_heapify(a, len(a), 0)  # @line L5
    return top  # @line L6

test_heap_extract_max.py 53 lines

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/heap-extract-max/. The authors are listed in prov/4GYGY9.jsonl.

import heapq
import random

import pytest

from heap_extract_max import heap_extract_max, max_heapify


def is_max_heap(a: list[int]) -> bool:
    return all(a[(j - 1) // 2] >= a[j] for j in range(1, len(a)))


def build(keys: list[int]) -> list[int]:
    a = list(keys)
    for i in range(len(a) // 2 - 1, -1, -1):
        max_heapify(a, len(a), i)
    return a


def test_clrs_example() -> None:
    a = [16, 14, 10, 8, 7, 9, 3, 2, 4, 1]
    assert heap_extract_max(a) == 16
    assert a == [14, 8, 10, 4, 7, 9, 3, 2, 1]
    assert heap_extract_max(a) == 14
    assert a == [10, 8, 9, 4, 7, 1, 3, 2]


def test_edges() -> None:
    with pytest.raises(IndexError):
        heap_extract_max([])
    a = [5]
    assert heap_extract_max(a) == 5 and a == []
    a = [5, 5, 5]
    assert heap_extract_max(a) == 5 and a == [5, 5]
    a = [2, 1]
    assert heap_extract_max(a) == 2 and a == [1]


def test_random_against_heapq() -> None:
    rng = random.Random(64)
    for _ in range(800):
        keys = [rng.randint(-30, 30) for _ in range(rng.randint(1, 40))]
        a = build(keys)
        assert is_max_heap(a)
        model = [-k for k in keys]
        heapq.heapify(model)
        while a:
            assert heap_extract_max(a) == -heapq.heappop(model)
            assert is_max_heap(a)
            assert sorted(a) == sorted(-k for k in model)

What the next level needs

the tests did not run in this build

Read the full page about Heap extract-max