Algorithm Search⌘K

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
bst_deletion.pyRawGitHub
# 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

Read the full page about Binary search tree deletion