Binary search tree deletion in Python
Handles leaf, one-child and two-child cases, the last by replacing the key with its successor.
- File
bst_deletion.py, 94 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/bst-deletion/. The authors are listed in prov/RSMQKE.jsonl.
from typing import Generic, TypeVar
K = TypeVar("K")
V = TypeVar("V")
class Node(Generic[K, V]):
"""A node of a binary search tree, with a parent pointer."""
__slots__ = ("key", "value", "left", "right", "p")
def __init__(self, key: K, value: V = None) -> None: # type: ignore[assignment]
self.key = key
self.value = value
self.left: "Node[K, V] | None" = None
self.right: "Node[K, V] | None" = None
self.p: "Node[K, V] | None" = None
class Tree(Generic[K, V]):
"""A binary search tree: a root and a count of nodes."""
def __init__(self) -> None:
self.root: Node[K, V] | None = None
self.n = 0
def insert(self, k: K, v: V = None) -> None: # type: ignore[assignment]
"""Leaf insertion, used to build trees for deletion."""
y, x = None, self.root
while x is not None:
y = x
if k < x.key:
x = x.left
elif x.key < k:
x = x.right
else:
x.value = v
return
z = Node(k, v)
z.p = y
if y is None:
self.root = z
elif k < y.key:
y.left = z
else:
y.right = z
self.n += 1
def tree_delete(t: Tree[K, V], k: K) -> bool:
"""Remove key k and its value from the BST t. Return True if k was a key.
A node with two children takes the key and value of its successor, and
the successor node, which has no left child, is spliced out.
Precondition: t has the BST property. Cost: O(h + 1) for height h.
"""
z = t.root # @line L1
while z is not None and (k < z.key or z.key < k):
z = z.left if k < z.key else z.right
if z is None: # @line L2
return False
if z.left is not None and z.right is not None: # @line L3
y = z.right # @line L4
while y.left is not None:
y = y.left
z.key, z.value = y.key, y.value # @line L5
z = y # @line L6
c = z.left if z.left is not None else z.right # @line L7
if c is not None: # @line L8
c.p = z.p
if z.p is None: # @line L9
t.root = c
elif z is z.p.left: # @line L10
z.p.left = c
else:
z.p.right = c # @line L11
t.n -= 1 # @line L12
return True
def inorder(t: Tree[K, V]) -> list[K]:
"""The keys of t in in-order, with an explicit stack."""
out, stack, x = [], [], t.root
while stack or x is not None:
while x is not None:
stack.append(x)
x = x.left
x = stack.pop()
out.append(x.key)
x = x.right
return out
test_bst_deletion.py 90 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/bst-deletion/. The authors are listed in prov/RSMQKE.jsonl.
import random
from bst_deletion import Tree, inorder, tree_delete
def check(t: Tree) -> None:
count, stack = 0, []
if t.root is not None:
assert t.root.p is None
stack.append((t.root, None, None))
while stack:
x, lo, hi = stack.pop()
count += 1
assert (lo is None or lo < x.key) and (hi is None or x.key < hi)
for c, a, b in ((x.left, lo, x.key), (x.right, x.key, hi)):
if c is not None:
assert c.p is x
stack.append((c, a, b))
assert count == t.n
def shape(x) -> object:
return None if x is None else (x.key, shape(x.left), shape(x.right))
def test_example() -> None:
t: Tree[int, int] = Tree()
for k in [50, 30, 70, 20, 40, 60, 80, 35, 45, 65]:
t.insert(k, k)
assert tree_delete(t, 60)
check(t)
assert shape(t.root.right) == (70, (65, None, None), (80, None, None))
assert tree_delete(t, 30)
check(t)
assert shape(t.root.left) == (35, (20, None, None), (40, None, (45, None, None)))
assert tree_delete(t, 20)
check(t)
assert tree_delete(t, 50)
check(t)
assert t.root.key == 65 and t.root.value == 65
assert shape(t.root) == (65, (35, None, (40, None, (45, None, None))), (70, None, (80, None, None)))
assert not tree_delete(t, 50)
assert inorder(t) == [35, 40, 45, 65, 70, 80]
def test_edges() -> None:
t: Tree[int, None] = Tree()
assert not tree_delete(t, 1)
t.insert(1)
assert tree_delete(t, 1) and t.root is None and t.n == 0
for k in (2, 1, 3):
t.insert(k)
assert tree_delete(t, 2)
check(t)
assert t.root.key == 3 and t.root.left.key == 1
def test_chain() -> None:
t: Tree[int, None] = Tree()
for k in range(3000):
t.insert(k)
for k in range(0, 3000, 3):
assert tree_delete(t, k)
check(t)
assert inorder(t) == [k for k in range(3000) if k % 3]
def test_random_against_model() -> None:
rng = random.Random(6620)
for _ in range(500):
t: Tree[int, int] = Tree()
model: dict[int, int] = {}
for step in range(rng.randint(0, 50)):
k = rng.randint(0, 25)
if rng.random() < 0.55:
t.insert(k, step)
model[k] = step
else:
assert tree_delete(t, k) == (k in model)
model.pop(k, None)
check(t)
assert inorder(t) == sorted(model)
stack = [t.root] if t.root else []
while stack:
x = stack.pop()
assert x.value == model[x.key]
stack += [c for c in (x.left, x.right) if c is not None]
What the next level needs
- the tests did not run in this build