Binary search tree insertion in Python
Walk down to a missing child position and attach the new key as a leaf.
- File
bst_insertion.py, 80 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-insertion/. The authors are listed in prov/82JNAM.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 tree_insert(t: Tree[K, V], k: K, v: V = None) -> Node[K, V]: # type: ignore[assignment]
"""Insert key k with value v into the BST t and return its node.
If k is already a key, its value is replaced and no node is added.
Precondition: t has the BST property. The keys need only "<".
Cost: O(h + 1) comparisons, where h is the height of t. Extra space O(1).
"""
y: Node[K, V] | None = None # @line L1
x = t.root
while x is not None: # @line L2
y = x # @line L3
if k < x.key: # @line L4
x = x.left
elif x.key < k: # @line L5
x = x.right
else:
x.value = v # @line L6
return x
z = Node(k, v) # @line L7
z.p = y
if y is None: # @line L8
t.root = z
elif k < y.key: # @line L9
y.left = z
else:
y.right = z # @line L10
t.n += 1 # @line L11
return z
def depth(z: Node[K, V]) -> int:
"""The number of edges from the root to z."""
d = 0
while z.p is not None:
z = z.p
d += 1
return d
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_insertion.py 67 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/bst-insertion/. The authors are listed in prov/82JNAM.jsonl.
import random
from bst_insertion import Node, Tree, depth, inorder, tree_insert
def check(t: Tree) -> None:
"""The BST property, the parent links and the count."""
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 test_example() -> None:
t: Tree[int, None] = Tree()
keys = [50, 30, 70, 20, 40, 60, 80, 35, 45, 65]
depths = [depth(tree_insert(t, k)) for k in keys]
check(t)
assert depths == [0, 1, 1, 2, 2, 2, 2, 3, 3, 3]
assert sum(depths) == 19
z = tree_insert(t, 42)
assert z.p is not None and z.p.key == 45 and z is z.p.left and depth(z) == 4
assert inorder(t) == sorted(keys + [42])
def test_duplicate_replaces_value() -> None:
t: Tree[int, str] = Tree()
a = tree_insert(t, 5, "a")
b = tree_insert(t, 5, "b")
assert a is b and t.n == 1 and a.value == "b"
def test_sorted_input_cost() -> None:
t: Tree[int, None] = Tree()
n = 2000
total = sum(depth(tree_insert(t, k)) for k in range(n))
assert total == n * (n - 1) // 2
check(t)
def test_random_against_sorted() -> None:
rng = random.Random(9031)
for _ in range(500):
t: Tree[int, int] = Tree()
keys = [rng.randint(0, 50) for _ in range(rng.randint(0, 40))]
for i, k in enumerate(keys):
tree_insert(t, k, i)
check(t)
assert inorder(t) == sorted(set(keys))
last = {k: i for i, k in enumerate(keys)}
stack = [t.root] if t.root else []
while stack:
x: Node = stack.pop()
assert x.value == last[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