Binary search tree search in Python
Descend from the root, going left or right by key comparison; O(h) time.
- File
bst_search.py, 70 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-search/. The authors are listed in prov/EJR7Z3.jsonl.
from typing import Any, Generic, TypeVar
K = TypeVar("K")
V = TypeVar("V")
class Node(Generic[K, V]):
"""A node of a binary search tree."""
__slots__ = ("key", "value", "left", "right")
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
def tree_search(x: Node[K, V] | None, k: K) -> Node[K, V] | None:
"""Return the node with key k in the BST rooted at x, or None.
Precondition: the tree rooted at x has the BST property. The keys need
only "<"; two keys are equal when neither is less than the other.
Cost: O(d + 1) comparisons, where d is the depth where the search stops;
O(h + 1) in the worst case. Extra space: O(1).
"""
while x is not None: # @line L1
if not (k < x.key or x.key < k): # @line L2
return x
if k < x.key: # @line L3
x = x.left
else:
x = x.right # @line L4
return None # @line L5
def tree_search_path(x: Node[K, V] | None, k: K) -> list[Any]:
"""The keys compared by tree_search(x, k), in order."""
path = []
while x is not None:
path.append(x.key)
if not (k < x.key or x.key < k):
break
x = x.left if k < x.key else x.right
return path
def insert(root: Node[K, V] | None, k: K, v: V = None) -> Node[K, V]: # type: ignore[assignment]
"""Insert k into the BST rooted at root and return the root (a helper)."""
z = Node(k, v)
if root is None:
return z
x = root
while True:
if k < x.key:
if x.left is None:
x.left = z
return root
x = x.left
elif x.key < k:
if x.right is None:
x.right = z
return root
x = x.right
else:
x.value = v
return root
test_bst_search.py 50 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/bst-search/. The authors are listed in prov/EJR7Z3.jsonl.
import random
from bst_search import Node, insert, tree_search, tree_search_path
def build(keys: list[int]) -> Node[int, int] | None:
root = None
for k in keys:
root = insert(root, k, 10 * k)
return root
def test_example() -> None:
root = build([15, 6, 18, 3, 7, 17, 20, 2, 4, 13, 9])
x = tree_search(root, 13)
assert x is not None and x.value == 130
assert tree_search_path(root, 13) == [15, 6, 7, 13]
assert tree_search(root, 8) is None
assert tree_search_path(root, 8) == [15, 6, 7, 13, 9]
assert tree_search_path(root, 15) == [15]
def test_edges() -> None:
assert tree_search(None, 1) is None
root = build([5])
assert tree_search(root, 5) is root
assert tree_search(root, 4) is None and tree_search(root, 6) is None
def test_chain_is_linear() -> None:
root = build(list(range(5000)))
assert tree_search(root, 4999) is not None
assert len(tree_search_path(root, 4999)) == 5000
assert tree_search(root, 5000) is None
def test_random_against_set() -> None:
rng = random.Random(7351)
for _ in range(600):
keys = [rng.randint(-40, 40) for _ in range(rng.randint(0, 30))]
root = build(keys)
s = set(keys)
for q in range(-45, 46):
x = tree_search(root, q)
assert (x is not None) == (q in s)
if x is not None:
assert x.key == q and x.value == 10 * q
What the next level needs
- the tests did not run in this build