1"""
2Red-Black Tree - Self-Balancing Binary Search Tree
3
4A Red-Black tree is a self-balancing binary search tree where each node
5has an extra bit for color (red or black). The tree maintains balance
6through five properties:
7 1. Every node is either red or black.
8 2. The root is black.
9 3. Every leaf (NIL) is black.
10 4. If a node is red, both its children are black.
11 5. All paths from a node to descendant leaves have the same black-height.
12
13Time Complexity:
14 - add: O(log n)
15 - search: O(log n)
16 - delete: O(log n)
17 - find_successor: O(log n)
18 - traverse: O(n)
19
20Space Complexity:
21 - Storage: O(n)
22 - Operations: O(log n) for iterative, O(n) for traversals
23"""
24
25from __future__ import annotations
26from enum import Enum
27
28
29class Color(Enum):
30 RED = 1
31 BLACK = 0
32
33
34class RbNode:
35 """A node in the Red-Black tree with color tracking."""
36
37 __slots__ = ("val", "left", "right", "parent", "color")
38
39 def __init__(self, val=None, color: Color = Color.BLACK):
40 self.val = val
41 # Self-references are placeholders; overwritten for real nodes
42 self.left: RbNode = self
43 self.right: RbNode = self
44 self.parent: RbNode = self
45 self.color: Color = color
46
47 def __str__(self):
48 return str(self.val)
49
50 def __repr__(self):
51 return str(self.val)
52
53
54class RedBlackTree:
55 """
56 Self-balancing Red-Black tree implementation.
57
58 Maintains the Red-Black invariants through color flips and rotations.
59 Uses a NIL sentinel node for leaves to simplify boundary conditions.
60 Provides operations for adding, searching, deleting nodes,
61 finding successors, and traversing the tree.
62 """
63
64 def __init__(self):
65 """Initialize an empty Red-Black tree."""
66 # NIL Sentinel Pattern:
67 # Instead of using None for empty children/parent, we use a single
68 # shared sentinel node. This simplifies the code because:
69 # - No None checks needed: all pointers are valid RbNode objects
70 # - Can safely access .color, .parent on any node (NIL is black)
71 # - Leaf semantics: RB theory treats leaves as black NIL nodes
72 # The sentinel has val=None and points to itself for all references.
73 # Real nodes must never use None as their value.
74 self.NIL = RbNode(val=None, color=Color.BLACK)
75 self.NIL.left = self.NIL.right = self.NIL.parent = self.NIL
76
77 self.root: RbNode = self.NIL # Empty tree
78
79 # -------------------------------------------------------------------------
80 # Basic Helpers
81 # -------------------------------------------------------------------------
82
83 def __contains__(self, val):
84 """Allow 'val in tree' syntax."""
85 return self.search(val) is not None
86
87 # -------------------------------------------------------------------------
88 # Rotations
89 # -------------------------------------------------------------------------
90 # Rotations preserve BST ordering while rebalancing the tree.
91 #
92 # Left Rotation: Right Rotation:
93 # x y y x
94 # / \ => / \ / \ => / \
95 # A y x C x C A y
96 # / \ / \ / \ / \
97 # B C A B A B B C
98 # -------------------------------------------------------------------------
99
100 def _rotate_left(self, x: RbNode) -> None:
101 """Perform left rotation on node x."""
102 y = x.right
103 # Turn y's left subtree into x's right subtree
104 x.right = y.left
105 if y.left is not self.NIL:
106 y.left.parent = x
107 # Link x's parent to y
108 y.parent = x.parent
109 if x.parent is self.NIL:
110 self.root = y
111 elif x is x.parent.left:
112 x.parent.left = y
113 else:
114 x.parent.right = y
115 # Put x on y's left
116 y.left = x
117 x.parent = y
118
119 def _rotate_right(self, y: RbNode) -> None:
120 """Perform right rotation on node y."""
121 x = y.left
122 # Turn x's right subtree into y's left subtree
123 y.left = x.right
124 if x.right is not self.NIL:
125 x.right.parent = y
126 # Link y's parent to x
127 x.parent = y.parent
128 if y.parent is self.NIL:
129 self.root = x
130 elif y is y.parent.right:
131 y.parent.right = x
132 else:
133 y.parent.left = x
134 # Put y on x's right
135 x.right = y
136 y.parent = x
137
138 # -------------------------------------------------------------------------
139 # Insert
140 # -------------------------------------------------------------------------
141
142 def add(self, val) -> None:
143 """Add a value to the tree. Raises KeyError if duplicate."""
144 # Create new red node with NIL children/parent
145 node = RbNode(val=val, color=Color.RED)
146 node.left = node.right = node.parent = self.NIL
147
148 # Standard BST insertion
149 y = self.NIL # Trailing pointer (will be parent)
150 x = self.root
151 while x is not self.NIL:
152 y = x
153 if val < x.val:
154 x = x.left
155 elif val > x.val:
156 x = x.right
157 else:
158 raise KeyError("Node with this value already in tree.")
159
160 # Insert node as child of y
161 node.parent = y
162 if y is self.NIL:
163 self.root = node # Tree was empty
164 elif node.val < y.val:
165 y.left = node
166 else:
167 y.right = node
168
169 # Fix Red-Black properties
170 self._insert_fixup(node)
171
172 def _insert_fixup(self, z: RbNode) -> None:
173 """Restore Red-Black properties after insertion."""
174 # Loop until parent is black (no violation) or z is root
175 while z.parent.color == Color.RED:
176 if z.parent is z.parent.parent.right:
177 y = z.parent.parent.left # Uncle
178 if y.color == Color.RED:
179 # Case 1: Uncle is red - recolor and move up
180 y.color = Color.BLACK
181 z.parent.color = Color.BLACK
182 z.parent.parent.color = Color.RED
183 z = z.parent.parent
184 else:
185 # Uncle is black
186 if z is z.parent.left:
187 # Case 2: z is left child - rotate to make it right child
188 z = z.parent
189 self._rotate_right(z)
190 # Case 3: z is right child - rotate and recolor
191 z.parent.color = Color.BLACK
192 z.parent.parent.color = Color.RED
193 self._rotate_left(z.parent.parent)
194 else:
195 # Mirror of above (parent is left child)
196 y = z.parent.parent.right # Uncle
197 if y.color == Color.RED:
198 # Case 1: Uncle is red - recolor and move up
199 y.color = Color.BLACK
200 z.parent.color = Color.BLACK
201 z.parent.parent.color = Color.RED
202 z = z.parent.parent
203 else:
204 # Uncle is black
205 if z is z.parent.right:
206 # Case 2: z is right child - rotate to make it left child
207 z = z.parent
208 self._rotate_left(z)
209 # Case 3: z is left child - rotate and recolor
210 z.parent.color = Color.BLACK
211 z.parent.parent.color = Color.RED
212 self._rotate_right(z.parent.parent)
213 if z is self.root:
214 break
215 # Ensure root is always black
216 self.root.color = Color.BLACK
217
218 # -------------------------------------------------------------------------
219 # Search / Min / Max / Successor / Predecessor
220 # -------------------------------------------------------------------------
221
222 def search(self, val) -> RbNode | None:
223 """Search for a value. Returns the node if found, None otherwise."""
224 node = self.root
225 while node is not self.NIL: # Stop at NIL sentinel (not None)
226 if val == node.val:
227 return node
228 node = node.left if val < node.val else node.right
229 return None # Not found - return None to external callers
230
231 def find_successor(self, node: RbNode) -> RbNode | None:
232 """
233 Find the in-order successor of a node.
234 If node has right subtree: successor is leftmost in right subtree.
235 Otherwise: successor is first ancestor where node is in left subtree.
236 """
237 if node.right is not self.NIL:
238 # Find leftmost node in right subtree
239 leftmost = node.right
240 while leftmost.left is not self.NIL:
241 leftmost = leftmost.left
242 return leftmost
243
244 # Walk up until we find an ancestor from the left
245 y = node.parent
246 while y is not self.NIL and node is y.right:
247 node = y
248 y = y.parent
249 return None if y is self.NIL else y
250
251 # -------------------------------------------------------------------------
252 # Delete
253 # -------------------------------------------------------------------------
254
255 def delete(self, val) -> None:
256 """Delete a value from the tree. No-op if value not found."""
257 # Find node to delete
258 z = self.root
259 while z is not self.NIL and z.val != val:
260 z = z.left if val < z.val else z.right
261 if z is self.NIL:
262 return # Not found
263
264 y = z # Node to be removed or moved
265 y_original_color = y.color
266
267 if z.left is self.NIL:
268 # Case 1: No left child - replace with right
269 x = z.right
270 self._transplant(z, z.right)
271 elif z.right is self.NIL:
272 # Case 2: No right child - replace with left
273 x = z.left
274 self._transplant(z, z.left)
275 else:
276 # Case 3: Two children - replace with successor
277 # Find successor (minimum in right subtree)
278 y = z.right
279 while y.left is not self.NIL:
280 y = y.left
281 y_original_color = y.color
282 x = y.right # Successor's right child (may be NIL)
283
284 if y.parent is z:
285 # Successor is direct child of z
286 x.parent = y # Maintain parent even if x is NIL
287 else:
288 # Successor is deeper - move it up
289 self._transplant(y, y.right)
290 y.right = z.right
291 y.right.parent = y
292
293 # Replace z with successor y
294 self._transplant(z, y)
295 y.left = z.left
296 y.left.parent = y
297 y.color = z.color
298
299 # Fix Red-Black properties if black node was removed
300 if y_original_color == Color.BLACK:
301 self._delete_fixup(x)
302
303 def _transplant(self, u: RbNode, v: RbNode) -> None:
304 """Replace subtree rooted at u with subtree rooted at v."""
305 if u.parent is self.NIL:
306 self.root = v
307 elif u is u.parent.left:
308 u.parent.left = v
309 else:
310 u.parent.right = v
311 v.parent = u.parent
312
313 def _delete_fixup(self, x: RbNode) -> None:
314 """Restore Red-Black properties after deletion."""
315 # Loop until x is red or root (extra black absorbed)
316 while x is not self.root and x.color == Color.BLACK:
317 if x is x.parent.left:
318 w = x.parent.right # Sibling
319 if w.color == Color.RED:
320 # Case 1: Sibling is red - rotate to make sibling black
321 w.color = Color.BLACK
322 x.parent.color = Color.RED
323 self._rotate_left(x.parent)
324 w = x.parent.right # New sibling
325 if w.left.color == Color.BLACK and w.right.color == Color.BLACK:
326 # Case 2: Sibling and both nephews are black - recolor
327 w.color = Color.RED
328 x = x.parent # Move extra black up
329 else:
330 if w.right.color == Color.BLACK:
331 # Case 3: Sibling's right child is black - rotate sibling
332 w.left.color = Color.BLACK
333 w.color = Color.RED
334 self._rotate_right(w)
335 w = x.parent.right
336 # Case 4: Sibling's right child is red - rotate parent
337 w.color = x.parent.color
338 x.parent.color = Color.BLACK
339 w.right.color = Color.BLACK
340 self._rotate_left(x.parent)
341 x = self.root # Done
342 else:
343 # Mirror of above (x is right child)
344 w = x.parent.left # Sibling
345 if w.color == Color.RED:
346 # Case 1: Sibling is red
347 w.color = Color.BLACK
348 x.parent.color = Color.RED
349 self._rotate_right(x.parent)
350 w = x.parent.left
351 if w.right.color == Color.BLACK and w.left.color == Color.BLACK:
352 # Case 2: Sibling and both nephews are black
353 w.color = Color.RED
354 x = x.parent
355 else:
356 if w.left.color == Color.BLACK:
357 # Case 3: Sibling's left child is black
358 w.right.color = Color.BLACK
359 w.color = Color.RED
360 self._rotate_left(w)
361 w = x.parent.left
362 # Case 4: Sibling's left child is red
363 w.color = x.parent.color
364 x.parent.color = Color.BLACK
365 w.left.color = Color.BLACK
366 self._rotate_right(x.parent)
367 x = self.root
368 # Remove extra black by coloring x black
369 x.color = Color.BLACK
370
371 # -------------------------------------------------------------------------
372 # Traversal Methods
373 # -------------------------------------------------------------------------
374
375 class TraverseOrder(Enum):
376 """Enumeration of tree traversal orders."""
377
378 preorder = 1
379 inorder = 2
380 postorder = 3
381
382 def traverse(self, order: TraverseOrder) -> list:
383 """Traverse the tree in the specified order. Returns list of values."""
384 result = []
385
386 def _preorder(node: RbNode):
387 if node is self.NIL:
388 return
389 result.append(node.val)
390 _preorder(node.left)
391 _preorder(node.right)
392
393 def _inorder(node: RbNode):
394 if node is self.NIL:
395 return
396 _inorder(node.left)
397 result.append(node.val)
398 _inorder(node.right)
399
400 def _postorder(node: RbNode):
401 if node is self.NIL:
402 return
403 _postorder(node.left)
404 _postorder(node.right)
405 result.append(node.val)
406
407 traversals = {
408 RedBlackTree.TraverseOrder.preorder: _preorder,
409 RedBlackTree.TraverseOrder.inorder: _inorder,
410 RedBlackTree.TraverseOrder.postorder: _postorder,
411 }
412
413 if order not in traversals:
414 raise TypeError("Invalid traverse order passed.")
415
416 traversals[order](self.root)
417 return result
418
419 def inorder_traversal(self) -> list:
420 """Return values in sorted (in-order) sequence."""
421 return self.traverse(RedBlackTree.TraverseOrder.inorder)
422
423 def levels_traversal(self) -> dict[int, list]:
424 """Return values grouped by level (breadth-first). Level 0 is root."""
425 if self.root is self.NIL:
426 return {}
427
428 levels = {}
429 queue = [(self.root, 0)]
430
431 for node, level in queue:
432 if level not in levels:
433 levels[level] = []
434 levels[level].append(node.val)
435
436 if node.left is not self.NIL:
437 queue.append((node.left, level + 1))
438 if node.right is not self.NIL:
439 queue.append((node.right, level + 1))
440
441 return levels
442
443 # -------------------------------------------------------------------------
444 # Depth Helpers
445 # -------------------------------------------------------------------------
446
447 def min_depth(self) -> int:
448 """Return minimum depth (shortest path from root to a leaf). -1 if empty."""
449 if self.root is self.NIL:
450 return -1
451
452 min_lvl = -1
453
454 def _dfs(node: RbNode, depth: int):
455 nonlocal min_lvl
456 if node is self.NIL:
457 return
458 # Found a leaf (both children are NIL)
459 if node.left is self.NIL and node.right is self.NIL:
460 if min_lvl == -1 or depth < min_lvl:
461 min_lvl = depth
462 _dfs(node.left, depth + 1)
463 _dfs(node.right, depth + 1)
464
465 _dfs(self.root, 0)
466 return min_lvl
467
468 def max_depth(self) -> int:
469 """Return maximum depth (longest path from root to a leaf). -1 if empty."""
470 if self.root is self.NIL:
471 return -1
472
473 max_lvl = 0
474
475 def _dfs(node: RbNode, depth: int):
476 nonlocal max_lvl
477 if node is self.NIL:
478 return
479 if depth > max_lvl:
480 max_lvl = depth
481 _dfs(node.left, depth + 1)
482 _dfs(node.right, depth + 1)
483
484 _dfs(self.root, 0)
485 return max_lvl
486
487
488# Example usage:
489if __name__ == "__main__":
490 tree = RedBlackTree()
491
492 # Insert values
493 for val in [10, 20, 30, 40, 50, 25]:
494 tree.add(val)
495
496 # Traverse in different orders
497 print("Inorder: ", tree.traverse(RedBlackTree.TraverseOrder.inorder)) # [10, 20, 25, 30, 40, 50]
498 print("Preorder: ", tree.traverse(RedBlackTree.TraverseOrder.preorder))
499 print("Postorder:", tree.traverse(RedBlackTree.TraverseOrder.postorder))
500 print("By levels:", tree.levels_traversal())
501
502 # Search
503 print("Search 25:", tree.search(25)) # RbNode with val=25
504 print("25 in tree:", 25 in tree) # True
505
506 # Delete
507 tree.delete(40)
508 tree.delete(30)
509 print("After deletions:", tree.inorder_traversal())
510