Red-black tree in Python
A BST with red/black node colours obeying four invariants that bound the height by 2 log2(n+1); the standard balanced BST in libraries.
- File
red_black_tree.py, 307 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/red-black-tree/. The authors are listed in prov/WQME6V.jsonl.
from __future__ import annotations
from typing import Any, Generic, Protocol, TypeVar
class Ordered(Protocol):
def __lt__(self, other: "Ordered", /) -> bool: ...
K = TypeVar("K", bound=Ordered)
RED, BLACK = True, False
class Node:
__slots__ = ("key", "color", "left", "right", "p")
def __init__(self, key: Any, color: bool, nil: "Node | None") -> None:
self.key = key
self.color = color
self.left: Node = nil if nil is not None else self
self.right: Node = nil if nil is not None else self
self.p: Node = nil if nil is not None else self
class RedBlackTree(Generic[K]):
"""An ordered set of keys kept in a red-black tree (CLRS style, with a sentinel).
contains, insert and delete take O(log n) time in the worst case for n keys.
An insertion does at most 2 rotations and a deletion at most 3. Keys must be
totally ordered by <; equal keys are stored once. No recursion is used.
"""
def __init__(self) -> None:
self.nil = Node(None, BLACK, None)
self.root: Node = self.nil
self.n = 0
self.rotations = 0
def __len__(self) -> int:
return self.n
# -- rotations ---------------------------------------------------------
def _left_rotate(self, x: Node) -> None:
self.rotations += 1
y = x.right
x.right = y.left
if y.left is not self.nil:
y.left.p = x
y.p = x.p
if x.p is self.nil:
self.root = y
elif x is x.p.left:
x.p.left = y
else:
x.p.right = y
y.left = x
x.p = y
def _right_rotate(self, y: Node) -> None:
self.rotations += 1
x = y.left
y.left = x.right
if x.right is not self.nil:
x.right.p = y
x.p = y.p
if y.p is self.nil:
self.root = x
elif y is y.p.left:
y.p.left = x
else:
y.p.right = x
x.right = y
y.p = x
# -- search ------------------------------------------------------------
def _find(self, k: K) -> Node:
x = self.root
while x is not self.nil:
if k < x.key:
x = x.left
elif x.key < k:
x = x.right
else:
return x
return self.nil
def contains(self, k: K) -> bool:
return self._find(k) is not self.nil
# -- insertion ---------------------------------------------------------
def insert(self, k: K) -> bool:
"""Insert k. Return True if k was not present."""
y, x = self.nil, self.root # @line L1
while x is not self.nil: # @line L2
y = x # @line L3
if k < x.key: # @line L4
x = x.left
elif x.key < k:
x = x.right
else:
return False
z = Node(k, RED, self.nil) # @line L5
z.p = y
if y is self.nil: # @line L6
self.root = z
elif k < y.key: # @line L7
y.left = z
else:
y.right = z
self.n += 1
self._insert_fixup(z) # @line L8
return True
def _insert_fixup(self, z: Node) -> None:
while z.p.color == RED: # @line L9
if z.p is z.p.p.left: # @line L10
u = z.p.p.right # @line L11
if u.color == RED: # @line L12
z.p.color = BLACK # @line L13
u.color = BLACK
z.p.p.color = RED
z = z.p.p
else: # @line L14
if z is z.p.right: # @line L15
z = z.p
self._left_rotate(z)
z.p.color = BLACK # @line L16
z.p.p.color = RED
self._right_rotate(z.p.p) # @line L17
else: # @line L18
u = z.p.p.left
if u.color == RED:
z.p.color = BLACK
u.color = BLACK
z.p.p.color = RED
z = z.p.p
else:
if z is z.p.left:
z = z.p
self._right_rotate(z)
z.p.color = BLACK
z.p.p.color = RED
self._left_rotate(z.p.p)
self.root.color = BLACK # @line L19
# -- deletion ----------------------------------------------------------
def _transplant(self, u: Node, v: Node) -> None:
if u.p is self.nil:
self.root = v
elif u is u.p.left:
u.p.left = v
else:
u.p.right = v
v.p = u.p
def delete(self, k: K) -> bool:
"""Delete k. Return True if k was present."""
z = self._find(k)
if z is self.nil:
return False
y = z
y_color = y.color
if z.left is self.nil:
x = z.right
self._transplant(z, z.right)
elif z.right is self.nil:
x = z.left
self._transplant(z, z.left)
else:
y = z.right
while y.left is not self.nil:
y = y.left
y_color = y.color
x = y.right
if y is not z.right:
self._transplant(y, y.right)
y.right = z.right
y.right.p = y
else:
x.p = y
self._transplant(z, y)
y.left = z.left
y.left.p = y
y.color = z.color
self.n -= 1
if y_color == BLACK:
self._delete_fixup(x)
return True
def _delete_fixup(self, x: Node) -> None:
while x is not self.root and x.color == BLACK:
if x is x.p.left:
w = x.p.right
if w.color == RED: # case 1
w.color = BLACK
x.p.color = RED
self._left_rotate(x.p)
w = x.p.right
if w.left.color == BLACK and w.right.color == BLACK: # case 2
w.color = RED
x = x.p
else:
if w.right.color == BLACK: # case 3
w.left.color = BLACK
w.color = RED
self._right_rotate(w)
w = x.p.right
w.color = x.p.color # case 4
x.p.color = BLACK
w.right.color = BLACK
self._left_rotate(x.p)
x = self.root
else:
w = x.p.left
if w.color == RED:
w.color = BLACK
x.p.color = RED
self._right_rotate(x.p)
w = x.p.left
if w.right.color == BLACK and w.left.color == BLACK:
w.color = RED
x = x.p
else:
if w.left.color == BLACK:
w.right.color = BLACK
w.color = RED
self._left_rotate(w)
w = x.p.left
w.color = x.p.color
x.p.color = BLACK
w.left.color = BLACK
self._right_rotate(x.p)
x = self.root
x.color = BLACK
# -- inspection --------------------------------------------------------
def keys(self) -> list[K]:
out: list[K] = []
stack: list[Node] = []
x = self.root
while stack or x is not self.nil:
while x is not self.nil:
stack.append(x)
x = x.left
x = stack.pop()
out.append(x.key)
x = x.right
return out
def height(self) -> int:
"""The number of nodes on the longest path from the root down."""
best = 0
stack = [(self.root, 0)]
while stack:
x, d = stack.pop()
if x is self.nil:
best = max(best, d)
else:
stack.append((x.left, d + 1))
stack.append((x.right, d + 1))
return best
def black_height(self) -> int:
"""The number of black nodes below the root on a path to a leaf."""
bh, x = 0, self.root
while x is not self.nil:
x = x.left
if x.color == BLACK:
bh += 1
return bh
def check_invariant(self) -> None:
"""Raise AssertionError unless the four red-black properties and the order hold."""
assert self.nil.color == BLACK
assert self.root.color == BLACK
assert self.root is self.nil or self.root.p is self.nil
count = 0
# iterative post-order: returns the black-height of each subtree
stack: list[tuple[Node, Any, Any, bool]] = [(self.root, None, None, False)]
bh: dict[int, int] = {}
while stack:
x, lo, hi, done = stack.pop()
if x is self.nil:
continue
if not done:
assert lo is None or lo < x.key
assert hi is None or x.key < hi
if x.color == RED:
assert x.left.color == BLACK and x.right.color == BLACK, "red node with a red child"
for c in (x.left, x.right):
assert c is self.nil or c.p is x, "bad parent pointer"
stack.append((x, lo, hi, True))
stack.append((x.left, lo, x.key, False))
stack.append((x.right, x.key, hi, False))
else:
count += 1
sides = []
for c in (x.left, x.right):
b = 0 if c is self.nil else bh[id(c)]
sides.append(b + (1 if c.color == BLACK else 0))
assert sides[0] == sides[1], "unequal black-heights"
bh[id(x)] = sides[0]
assert count == self.n
test_red_black_tree.py 93 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/red-black-tree/. The authors are listed in prov/WQME6V.jsonl.
import math
import random
from red_black_tree import RED, Node, RedBlackTree
def shape(t: RedBlackTree[int], x: "Node | None" = None) -> object:
x = t.root if x is None else x
if x is t.nil:
return None
me = f"{x.key}{'r' if x.color == RED else 'b'}"
if x.left is t.nil and x.right is t.nil:
return me
return (shape(t, x.left) if x.left is not t.nil else None, me,
shape(t, x.right) if x.right is not t.nil else None)
def test_page_example() -> None:
t: RedBlackTree[int] = RedBlackTree()
got = []
for k in [10, 20, 30, 15, 25, 5, 1]:
t.insert(k)
t.check_invariant()
got.append(shape(t))
assert got[2] == ("10r", "20b", "30r")
assert got[3] == ((None, "10b", "15r"), "20b", "30b")
assert got[-1] == ((("1r", "5b", None), "10r", "15b"), "20b", ("25r", "30b", None))
for k, want in [
(20, ((("1r", "5b", None), "10r", "15b"), "25b", "30b")),
(1, (("5b", "10r", "15b"), "25b", "30b")),
(30, ("5b", "10b", ("15r", "25b", None))),
]:
t.delete(k)
t.check_invariant()
assert shape(t) == want
def test_sorted_input() -> None:
t: RedBlackTree[int] = RedBlackTree()
for k in range(1, 1001):
t.insert(k)
t.check_invariant()
assert t.height() <= 2 * math.log2(1001)
for k in range(1, 1001, 2):
t.delete(k)
t.check_invariant()
assert t.keys() == list(range(2, 1001, 2))
def test_edges() -> None:
t: RedBlackTree[int] = RedBlackTree()
assert not t.delete(3) and not t.contains(3) and t.height() == 0
assert t.insert(3) and not t.insert(3) and len(t) == 1
assert t.delete(3) and len(t) == 0 and t.root is t.nil
t.check_invariant()
def test_rotation_counts() -> None:
rng = random.Random(1978)
t: RedBlackTree[int] = RedBlackTree()
for _ in range(4000):
k = rng.randint(0, 3000)
r0 = t.rotations
if rng.random() < 0.6:
t.insert(k)
assert t.rotations - r0 <= 2
else:
t.delete(k)
assert t.rotations - r0 <= 3
t.check_invariant()
def test_random_against_set() -> None:
rng = random.Random(1972)
for _ in range(500):
t: RedBlackTree[int] = RedBlackTree()
model: set[int] = set()
for _ in range(rng.randint(0, 80)):
k = rng.randint(0, 40)
if rng.random() < 0.6:
assert t.insert(k) == (k not in model)
model.add(k)
else:
assert t.delete(k) == (k in model)
model.discard(k)
assert t.contains(k) == (k in model)
t.check_invariant()
assert t.keys() == sorted(model)
n = len(model)
assert t.height() <= 2 * math.log2(n + 1)
What the next level needs
- the tests did not run in this build