495 lines
17 KiB
Python
495 lines
17 KiB
Python
import traceback
|
|
|
|
class Node:
|
|
def __init__(self, value, left=None, right=None, parent=None, is_red=False):
|
|
self.value = value
|
|
self.left: Node | None = left
|
|
self.right: Node | None = right
|
|
self.is_red = is_red
|
|
|
|
# required to traverse the tree and perform various checks
|
|
self.parent: Node | None = parent
|
|
|
|
class RedBlackTree:
|
|
def __init__(self):
|
|
self.root: Node | None = None
|
|
|
|
def insert(self, value):
|
|
# assumption: inserting the same value twice will
|
|
# insert nothing the second time
|
|
|
|
node = Node(value)
|
|
if self.root is None:
|
|
# the root is always black (Node defaults to black)
|
|
self.root = node
|
|
return
|
|
|
|
# insert like a normal binary search tree
|
|
prev: Node | None = None
|
|
current = self.root
|
|
while current:
|
|
if current.value == value:
|
|
return
|
|
prev = current
|
|
if current.value > value:
|
|
current = current.left
|
|
elif current.value < value:
|
|
current = current.right
|
|
|
|
# a newly inserted (non-root) node is always red, and it needs a
|
|
# parent pointer so rebalancing can walk back up toward the root
|
|
node.is_red = True
|
|
node.parent = prev
|
|
if prev and prev.value > value:
|
|
prev.left = node
|
|
elif prev and prev.value < value:
|
|
prev.right = node
|
|
|
|
self.rebalance_from_just_inserted(node)
|
|
|
|
def delete(self, value, root=None, parent=None):
|
|
# does nothing if the value doesn't exist
|
|
if parent is None:
|
|
root = self.root
|
|
if root is None:
|
|
return
|
|
|
|
if root.value == value:
|
|
new_root: Node | None
|
|
if root.left is None:
|
|
new_root = root.right
|
|
elif root.right is None:
|
|
new_root = root.left
|
|
else:
|
|
# both left and right subtrees exist
|
|
# make left subtree root the tree root, and re-attach
|
|
# the right subtree
|
|
right_subtree_root = root.right
|
|
new_root = root.left
|
|
|
|
prev = None
|
|
current = new_root
|
|
while current:
|
|
prev = current
|
|
current = current.right
|
|
prev.right = right_subtree_root
|
|
|
|
if parent is None:
|
|
self.root = new_root
|
|
else:
|
|
if parent.value > root.value:
|
|
parent.left = new_root
|
|
else:
|
|
parent.right = new_root
|
|
elif root.value > value:
|
|
self.delete(value, root.left, root)
|
|
elif root.value < value:
|
|
self.delete(value, root.right, root)
|
|
|
|
def rebalance_from_just_inserted(self, node: Node | None):
|
|
if node is None:
|
|
return
|
|
|
|
parent = node.parent
|
|
if parent is not None and not parent.is_red:
|
|
return
|
|
elif parent is None:
|
|
node.is_red = False
|
|
return
|
|
|
|
grandparent = parent.parent
|
|
if grandparent is None:
|
|
return
|
|
|
|
uncle = grandparent.right if grandparent.left == parent else grandparent.left
|
|
|
|
if uncle is not None and uncle.is_red:
|
|
uncle.is_red = False
|
|
parent.is_red = False
|
|
grandparent.is_red = True
|
|
|
|
self.rebalance_from_just_inserted(grandparent)
|
|
return
|
|
|
|
if grandparent.left == parent and parent.right == node:
|
|
node, parent = parent, self.rotate_left(parent)
|
|
elif grandparent.right == parent and parent.left == node:
|
|
node, parent = parent, self.rotate_right(parent)
|
|
|
|
new_grandparent: Node | None = None
|
|
if grandparent.left == parent and parent.left == node:
|
|
new_grandparent = self.rotate_right(grandparent)
|
|
elif grandparent.right == parent and parent.right == node:
|
|
new_grandparent = self.rotate_left(grandparent)
|
|
|
|
if grandparent == self.root and new_grandparent is not None:
|
|
self.root = new_grandparent
|
|
self.root.is_red = False
|
|
grandparent.is_red = True
|
|
|
|
|
|
def rotate_left(self, node: Node) -> Node:
|
|
if node.right is None:
|
|
return node
|
|
|
|
original_parent = node.parent
|
|
new_parent = node.right
|
|
new_right_child = node.right.left
|
|
|
|
node.right = new_right_child
|
|
if new_right_child is not None:
|
|
new_right_child.parent = node
|
|
|
|
new_parent.left = node
|
|
node.parent = new_parent
|
|
new_parent.parent = original_parent
|
|
|
|
if original_parent is not None:
|
|
if original_parent.left == node:
|
|
original_parent.left = new_parent
|
|
elif original_parent.right == node:
|
|
original_parent.right = new_parent
|
|
|
|
return new_parent
|
|
|
|
def rotate_right(self, node: Node) -> Node:
|
|
if node.left is None:
|
|
return node
|
|
|
|
original_parent = node.parent
|
|
new_parent = node.left
|
|
new_left_child = node.left.right
|
|
|
|
node.left = new_left_child
|
|
if new_left_child is not None:
|
|
new_left_child.parent = node
|
|
|
|
new_parent.right = node
|
|
node.parent = new_parent
|
|
new_parent.parent = original_parent
|
|
|
|
if original_parent is not None:
|
|
if original_parent.left == node:
|
|
original_parent.left = new_parent
|
|
elif original_parent.right == node:
|
|
original_parent.right = new_parent
|
|
|
|
return new_parent
|
|
|
|
def visualize(self) -> str:
|
|
# renders the tree top-down with the root on top and branches
|
|
# ( / and \ ) drawn down to each child. spacing is computed so that
|
|
# subtrees never overlap, no matter the shape of the tree.
|
|
if self.root is None:
|
|
return "<empty tree>"
|
|
|
|
# render() returns, for the subtree rooted at `node`:
|
|
# lines - the block of text drawing the subtree
|
|
# width - how many characters wide that block is
|
|
# height - how many lines tall that block is
|
|
# middle - the column where this node's value is centred (so the
|
|
# caller knows where to attach its branch)
|
|
def render(node: Node) -> tuple[list[str], int, int, int]:
|
|
# label shows the value plus its color: R (red) or B (black)
|
|
label = f"{node.value}{'R' if node.is_red else 'B'}"
|
|
label_width = len(label)
|
|
|
|
# leaf: just the value on a single line
|
|
if node.left is None and node.right is None:
|
|
return [label], label_width, 1, label_width // 2
|
|
|
|
# only a left child
|
|
if node.right is None:
|
|
lines, width, height, mid = render(node.left)
|
|
first = (mid + 1) * " " + (width - mid - 1) * "_" + label
|
|
second = mid * " " + "/" + (width - mid - 1 + label_width) * " "
|
|
shifted = [line + label_width * " " for line in lines]
|
|
return [first, second] + shifted, width + label_width, height + 2, width + label_width // 2
|
|
|
|
# only a right child
|
|
if node.left is None:
|
|
lines, width, height, mid = render(node.right)
|
|
first = label + mid * "_" + (width - mid) * " "
|
|
second = (label_width + mid) * " " + "\\" + (width - mid - 1) * " "
|
|
shifted = [label_width * " " + line for line in lines]
|
|
return [first, second] + shifted, width + label_width, height + 2, label_width // 2
|
|
|
|
# two children: render each side, then place this node between them
|
|
left_lines, left_w, left_h, left_mid = render(node.left)
|
|
right_lines, right_w, right_h, right_mid = render(node.right)
|
|
first = (
|
|
(left_mid + 1) * " "
|
|
+ (left_w - left_mid - 1) * "_"
|
|
+ label
|
|
+ right_mid * "_"
|
|
+ (right_w - right_mid) * " "
|
|
)
|
|
second = (
|
|
left_mid * " "
|
|
+ "/"
|
|
+ (left_w - left_mid - 1 + label_width + right_mid) * " "
|
|
+ "\\"
|
|
+ (right_w - right_mid - 1) * " "
|
|
)
|
|
# pad the shorter side so the two blocks line up row-for-row
|
|
if left_h < right_h:
|
|
left_lines += [left_w * " "] * (right_h - left_h)
|
|
elif right_h < left_h:
|
|
right_lines += [right_w * " "] * (left_h - right_h)
|
|
merged = [l + label_width * " " + r for l, r in zip(left_lines, right_lines)]
|
|
return (
|
|
[first, second] + merged,
|
|
left_w + right_w + label_width,
|
|
max(left_h, right_h) + 2,
|
|
left_w + label_width // 2,
|
|
)
|
|
|
|
lines, _, _, _ = render(self.root)
|
|
return "\n".join(lines)
|
|
|
|
|
|
# =============== TEST CASES ===================
|
|
#
|
|
# These run automatically: python red_black_tree.py
|
|
#
|
|
# Each case inserts ONE value into a hand-built red-black tree and checks the
|
|
# result, printing the tree BEFORE, the ACTUAL tree after your rebalance, and
|
|
# the EXPECTED tree. A case passes when the actual tree matches the expected
|
|
# tree AND the result is still a valid red-black tree.
|
|
#
|
|
# HOW TO ADD YOUR OWN CASE:
|
|
# Append a Case(...) to the CASES list below with four fields:
|
|
# name - short description
|
|
# before - a function returning the tree BEFORE the insert (None = empty)
|
|
# insert - the value to insert
|
|
# expected - a function returning the tree you expect AFTERWARD
|
|
# Build trees with node(value, "R" or "B", left=..., right=...). Any child
|
|
# you leave out is treated as a (black) nil leaf.
|
|
|
|
from dataclasses import dataclass
|
|
from typing import Callable
|
|
|
|
|
|
def node(value, color, left=None, right=None):
|
|
"""Build a Node with an explicit color ("R"/"B"), wiring parent pointers."""
|
|
n = Node(value, is_red=(color == "R"))
|
|
n.left = left
|
|
n.right = right
|
|
if left is not None:
|
|
left.parent = n
|
|
if right is not None:
|
|
right.parent = n
|
|
return n
|
|
|
|
|
|
def tree(root):
|
|
"""Wrap a root Node (or None) in a RedBlackTree."""
|
|
t = RedBlackTree()
|
|
t.root = root
|
|
return t
|
|
|
|
|
|
@dataclass
|
|
class Case:
|
|
name: str
|
|
before: Callable # () -> Node | None
|
|
insert: int
|
|
expected: Callable # () -> Node | None
|
|
|
|
|
|
CASES = [
|
|
Case(
|
|
"Empty tree -> black root",
|
|
before=lambda: None,
|
|
insert=25,
|
|
expected=lambda: node(25, "B"),
|
|
),
|
|
Case(
|
|
"Black parent -> new red child, no rebalancing",
|
|
before=lambda: node(25, "B"),
|
|
insert=15,
|
|
expected=lambda: node(25, "B", left=node(15, "R")),
|
|
),
|
|
Case(
|
|
"Duplicate insert -> tree unchanged",
|
|
before=lambda: node(25, "B", left=node(15, "R")),
|
|
insert=15,
|
|
expected=lambda: node(25, "B", left=node(15, "R")),
|
|
),
|
|
Case(
|
|
"Red uncle -> recolor (grandparent is the root)",
|
|
before=lambda: node(25, "B", left=node(15, "R"), right=node(35, "R")),
|
|
insert=10,
|
|
expected=lambda: node(
|
|
25, "B",
|
|
left=node(15, "B", left=node(10, "R")),
|
|
right=node(35, "B"),
|
|
),
|
|
),
|
|
Case(
|
|
"LL -> right-rotate the grandparent",
|
|
before=lambda: node(30, "B", left=node(20, "R")),
|
|
insert=10,
|
|
expected=lambda: node(20, "B", left=node(10, "R"), right=node(30, "R")),
|
|
),
|
|
Case(
|
|
"RR -> left-rotate the grandparent",
|
|
before=lambda: node(30, "B", right=node(40, "R")),
|
|
insert=50,
|
|
expected=lambda: node(40, "B", left=node(30, "R"), right=node(50, "R")),
|
|
),
|
|
Case(
|
|
"LR -> left-rotate parent, then right-rotate grandparent",
|
|
before=lambda: node(30, "B", left=node(20, "R")),
|
|
insert=25,
|
|
expected=lambda: node(25, "B", left=node(20, "R"), right=node(30, "R")),
|
|
),
|
|
Case(
|
|
"RL -> right-rotate parent, then left-rotate grandparent",
|
|
before=lambda: node(30, "B", right=node(40, "R")),
|
|
insert=35,
|
|
expected=lambda: node(35, "B", left=node(30, "R"), right=node(40, "R")),
|
|
),
|
|
Case(
|
|
"Cascade: recolor propagates up, then rotate near the root",
|
|
before=lambda: node(
|
|
11, "B",
|
|
left=node(
|
|
2, "R",
|
|
left=node(1, "B"),
|
|
right=node(7, "B", left=node(5, "R"), right=node(8, "R")),
|
|
),
|
|
right=node(14, "B", right=node(15, "R")),
|
|
),
|
|
insert=4,
|
|
expected=lambda: node(
|
|
7, "B",
|
|
left=node(
|
|
2, "R",
|
|
left=node(1, "B"),
|
|
right=node(5, "B", left=node(4, "R")),
|
|
),
|
|
right=node(
|
|
11, "R",
|
|
left=node(8, "B"),
|
|
right=node(14, "B", right=node(15, "R")),
|
|
),
|
|
),
|
|
),
|
|
]
|
|
|
|
|
|
def validate_rb(t):
|
|
"""Return a list of red-black property violations ([] means valid)."""
|
|
problems = []
|
|
root = t.root
|
|
if root is None:
|
|
return problems
|
|
if root.is_red:
|
|
problems.append("root is red")
|
|
if root.parent is not None:
|
|
problems.append("root has a non-None parent pointer")
|
|
|
|
black_heights = set()
|
|
seen = set()
|
|
|
|
def check(n, low, high, black_count):
|
|
if n is None:
|
|
black_heights.add(black_count + 1) # nil leaves count as black
|
|
return
|
|
if id(n) in seen:
|
|
# a proper tree never reaches the same node twice
|
|
problems.append(f"not a tree: node {n.value} reached twice (cycle or shared subtree)")
|
|
return
|
|
seen.add(id(n))
|
|
if low is not None and n.value <= low:
|
|
problems.append(f"BST order broken at {n.value}")
|
|
if high is not None and n.value >= high:
|
|
problems.append(f"BST order broken at {n.value}")
|
|
if n.is_red and ((n.left and n.left.is_red) or (n.right and n.right.is_red)):
|
|
problems.append(f"red node {n.value} has a red child")
|
|
if n.left is not None and n.left.parent is not n:
|
|
problems.append(f"broken parent pointer: left child {n.left.value} does not point back to {n.value}")
|
|
if n.right is not None and n.right.parent is not n:
|
|
problems.append(f"broken parent pointer: right child {n.right.value} does not point back to {n.value}")
|
|
nb = black_count + (0 if n.is_red else 1)
|
|
check(n.left, low, n.value, nb)
|
|
check(n.right, n.value, high, nb)
|
|
|
|
check(root, None, None, 0)
|
|
if len(black_heights) > 1:
|
|
problems.append(f"unequal black-heights on paths to nil: {sorted(black_heights)}")
|
|
return problems
|
|
|
|
|
|
def _indent(text: str) -> str:
|
|
return "\n".join(" " + line for line in text.split("\n"))
|
|
|
|
|
|
def run_tests():
|
|
passed = 0
|
|
for i, case in enumerate(CASES, 1):
|
|
print("=" * 64)
|
|
print(f"CASE {i}: {case.name}")
|
|
print("=" * 64)
|
|
|
|
t = tree(case.before())
|
|
print("BEFORE:")
|
|
print(_indent(t.visualize()))
|
|
print(f"\n insert({case.insert})\n")
|
|
|
|
# Everything the student's code can affect is inside this guard, so a
|
|
# buggy rebalance (even one that builds a cyclic/broken tree) fails just
|
|
# this case instead of aborting the whole suite.
|
|
try:
|
|
t.insert(case.insert)
|
|
actual = t.visualize()
|
|
expected_tree = tree(case.expected())
|
|
expected = expected_tree.visualize()
|
|
violations = validate_rb(t)
|
|
expected_problems = validate_rb(expected_tree)
|
|
except RecursionError:
|
|
print(" RESULT: ERROR - hit maximum recursion depth.")
|
|
print(" Your rebalance likely created a cycle or otherwise broke the")
|
|
print(" tree structure (a child pointing back up at an ancestor).\n")
|
|
continue
|
|
except Exception:
|
|
print(f" RESULT: ERROR - see stack trace below\n")
|
|
print(traceback.format_exc())
|
|
print("")
|
|
continue
|
|
|
|
# Guard the author (you) against a typo when adding a new case: the
|
|
# 'expected' tree should itself be a valid red-black tree.
|
|
if expected_problems:
|
|
print(" RESULT: BAD TEST - the 'expected' tree is not a valid red-black tree:")
|
|
for p in expected_problems:
|
|
print(f" - {p}")
|
|
print()
|
|
continue
|
|
|
|
print("ACTUAL (what your code produced):")
|
|
print(_indent(actual))
|
|
print("\nEXPECTED:")
|
|
print(_indent(expected))
|
|
print()
|
|
|
|
if actual == expected and not violations:
|
|
print(" RESULT: PASS")
|
|
passed += 1
|
|
else:
|
|
print(" RESULT: FAIL")
|
|
if actual != expected:
|
|
print(" - actual tree does not match the expected tree")
|
|
for v in violations:
|
|
print(f" - red-black property broken: {v}")
|
|
print()
|
|
|
|
print("=" * 64)
|
|
print(f"{passed}/{len(CASES)} cases passed")
|
|
print("=" * 64)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
run_tests()
|