BST floor and ceiling in Python
Queries on a binary search tree that find the largest key not above, or smallest key not below, a given key.
- File
bst_floor_and_ceiling.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-floor-and-ceiling/. The authors are listed in prov/85F3S3.jsonl.
from typing import Generic, TypeVar
K = TypeVar("K")
class Node(Generic[K]):
"""A node of a binary search tree."""
__slots__ = ("key", "left", "right")
def __init__(self, key: K) -> None:
self.key = key
self.left: "Node[K] | None" = None
self.right: "Node[K] | None" = None
def floor(x: Node[K] | None, k: K) -> K | None:
"""The largest key <= k in the BST rooted at x, or None.
Precondition: the tree has the BST property. Keys need only "<".
Cost: O(h + 1) comparisons for height h; O(1) extra space.
"""
best: Node[K] | None = None # @line L1
while x is not None: # @line L2
if not (k < x.key or x.key < k): # @line L3
return x.key
if k < x.key: # @line L4
x = x.left
else:
best = x # @line L5
x = x.right
return None if best is None else best.key # @line L6
def ceiling(x: Node[K] | None, k: K) -> K | None:
"""The smallest key >= k in the BST rooted at x, or None. Cost O(h + 1)."""
best: Node[K] | None = None # @line L7
while x is not None: # @line L8
if not (k < x.key or x.key < k): # @line L9
return x.key
if x.key < k: # @line L10
x = x.right
else:
best = x # @line L11
x = x.left
return None if best is None else best.key # @line L12
def insert(root: Node[K] | None, k: K) -> Node[K]:
"""Leaf insertion of k (a helper to build trees); returns the root."""
z = Node(k)
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:
return root
test_bst_floor_and_ceiling.py 51 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/bst-floor-and-ceiling/. The authors are listed in prov/85F3S3.jsonl.
import bisect
import random
from bst_floor_and_ceiling import ceiling, floor, insert
EXAMPLE = [50, 30, 70, 20, 40, 60, 80, 35, 45, 65]
def build(keys):
root = None
for k in keys:
root = insert(root, k)
return root
def test_example() -> None:
root = build(EXAMPLE)
assert floor(root, 62) == 60 and ceiling(root, 62) == 65
assert floor(root, 45) == 45 and ceiling(root, 45) == 45
assert floor(root, 10) is None and ceiling(root, 10) == 20
assert floor(root, 90) == 80 and ceiling(root, 90) is None
assert floor(root, 49) == 45 and ceiling(root, 49) == 50
assert floor(root, 36) == 35 and ceiling(root, 36) == 40
def test_edges() -> None:
assert floor(None, 1) is None and ceiling(None, 1) is None
root = build([5])
assert floor(root, 5) == 5 and floor(root, 4) is None and floor(root, 6) == 5
assert ceiling(root, 4) == 5 and ceiling(root, 6) is None
def test_chain() -> None:
root = build(list(range(0, 6000, 2)))
assert floor(root, 5999) == 5998 and ceiling(root, 1) == 2
def test_random_against_bisect() -> None:
rng = random.Random(3301)
for _ in range(600):
keys = [rng.randint(0, 80) for _ in range(rng.randint(0, 30))]
root = build(keys)
s = sorted(set(keys))
for q in range(-2, 84):
i = bisect.bisect_right(s, q)
assert floor(root, q) == (s[i - 1] if i else None)
j = bisect.bisect_left(s, q)
assert ceiling(root, q) == (s[j] if j < len(s) else None)
What the next level needs
- the tests did not run in this build