Heap delete in Python
A heap operation that removes an arbitrary element by replacing it with the last element and then restoring the heap property.
- File
heap_delete.py, 46 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/heap-delete/. The authors are listed in prov/8TNZV0.jsonl.
from typing import Any, Protocol, TypeVar
class Ordered(Protocol):
def __lt__(self, other: Any, /) -> bool: ...
T = TypeVar("T", bound=Ordered)
def sift_up(a: list[T], i: int) -> None:
while i > 0 and a[(i - 1) // 2] < a[i]: # @line L9
p = (i - 1) // 2
a[i], a[p] = a[p], a[i] # @line L10
i = p # @line L11
def max_heapify(a: list[T], n: int, i: int) -> None:
while True: # @line L12
left, right, m = 2 * i + 1, 2 * i + 2, i # @line L13
if left < n and a[m] < a[left]: # @line L14
m = left
if right < n and a[m] < a[right]: # @line L15
m = right
if m == i: # @line L16
return
a[i], a[m] = a[m], a[i] # @line L17
i = m
def heap_delete(a: list[T], i: int) -> T:
"""Remove and return a[i] from the max-heap a (0-based). O(log n) time."""
if not 0 <= i < len(a): # @line L1
raise IndexError("index out of range")
x = a[i] # @line L2
last = a.pop() # @line L3 @line L4
if i < len(a): # @line L5
a[i] = last
if i > 0 and a[(i - 1) // 2] < a[i]: # @line L6
sift_up(a, i)
else:
max_heapify(a, len(a), i) # @line L7
return x # @line L8
test_heap_delete.py 61 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/heap-delete/. The authors are listed in prov/8TNZV0.jsonl.
import random
import pytest
from heap_delete import heap_delete, 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_examples() -> None:
a = [16, 14, 10, 8, 7, 9, 3, 2, 4, 1]
assert heap_delete(a, 1) == 14
assert a == [16, 8, 10, 4, 7, 9, 3, 2, 1]
b = [100, 50, 90, 40, 45, 80, 85]
assert heap_delete(b, 3) == 40
assert b == [100, 85, 90, 50, 45, 80]
def test_edges() -> None:
with pytest.raises(IndexError):
heap_delete([], 0)
with pytest.raises(IndexError):
heap_delete([3, 2], 2)
a = [7]
assert heap_delete(a, 0) == 7 and a == []
a = [7, 5, 6]
assert heap_delete(a, 2) == 6 and a == [7, 5]
a = [7, 5, 6]
assert heap_delete(a, 0) == 7 and a == [6, 5]
def test_sift_down_only_is_wrong() -> None:
b = [100, 50, 90, 40, 45, 80, 85]
b[3] = b.pop()
max_heapify(b, len(b), 3)
assert not is_max_heap(b)
def test_random() -> None:
rng = random.Random(8)
for _ in range(1000):
a = build([rng.randint(-20, 20) for _ in range(rng.randint(1, 35))])
while a:
i = rng.randrange(len(a))
expect = sorted(a[:i] + a[i + 1 :])
x = a[i]
assert heap_delete(a, i) == x
assert is_max_heap(a)
assert sorted(a) == expect
What the next level needs
- the tests did not run in this build