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): # This is what you'll be implementing. # # By the time this is called, `node` has already been inserted like a # normal BST node, colored red, and had its `parent` pointer set (see # insert()). Your job is to restore the red-black properties by # rebalancing the subtree around `node`. # # Rebalancing is only needed when node.parent is red. Remember to keep # the root black at the end. pass 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 "" # 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 as e: print(f" RESULT: ERROR - {type(e).__name__}: {e}\n") 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()