"""
Red-Black Tree implementation for balanced binary search tree operations.
This module implements a self-balancing binary search tree using the Red-Black Tree
algorithm. The tree maintains logarithmic time complexity for insertion, deletion,
and search operations through color-based balancing rules.
Classes
-------
RedBlackTree
A self-balancing binary search tree that maintains Red-Black properties
to ensure O(log n) performance for all major operations.
"""
from typing import Callable, Sequence, Any
from node import *
from config import COLORS
[docs]
class RedBlackTree:
"""
A class to represent a Red-Black Tree.
Attributes
----------
_root: Node | None
The root of the tree.
_type: type
The type of the elements in the tree.
_key: Callable | None
A function to compare elements in the tree.
_smallest: Node | None
The smallest element in the tree.
_largest: Node | None
The largest element in the tree.
"""
def __init__(self, element_type: type, collection: Sequence[Any] | None = None, comparator: Callable[[object], Any] | None = None):
self._root: Node | None = None
self._type: type = element_type
self._key: Callable[[object], Any] | None = comparator
self._smallest: Node | None = None
self._largest: Node | None = None
if collection:
for element in collection:
self.insert(element)
@property
def root(self) -> Node | None:
"""
Return the root of the tree.
Returns:
Node | None: The root of the tree.
"""
return self._root
@property
def type(self) -> type:
"""
Return the type of the elements in the tree.
Returns:
type: The type of the elements in the tree.
"""
return self._type
@property
def smallest(self) -> object:
"""
Return the smallest element in the tree.
Returns:
object: The smallest element in the tree.
"""
if self._smallest:
return self._smallest.data
return None
@property
def largest(self) -> object:
"""
Return the largest element in the tree.
Returns:
object: The largest element in the tree.
"""
if self._largest:
return self._largest.data
return None
[docs]
def insert(self, value: object) -> tuple[int, Node] | tuple[int, Node|None]:
"""
Insert a value into the tree.
Args:
value (object): The value to insert.
Returns:
tuple[int, Node] | tuple[int, int]: (0, node) if inserted, (-1, None) otherwise.
Raises:
TypeError if value of unsupported type.
"""
if type(value) != self._type:
raise TypeError(f"Value must be of type {self._type}.")
if not self._root:
self._root = Node(value, 1, None, None, None)
self._largest = self._smallest = self._root
return 0, self._root
status, parent = self.find(value)
if status == 0:
return -1, None
if self.__compare(value, parent.data) > 0:
parent.right = Node(value, 0, parent, None, None)
node = parent.right
else:
parent.left = Node(value, 0, parent, None, None)
node = parent.left
if not self._smallest or self.__compare(node.data, self._smallest.data) < 0:
self._smallest = node
if not self._largest or self.__compare(node.data, self._largest.data) > 0:
self._largest = node
if parent != self._root:
self.__fix_insert(node)
return 0, node
[docs]
def remove(self, value: object) -> tuple[int, Node] | tuple[int, Node|None]:
"""
Remove a value from the tree.
Args:
value (object): The value to remove.
Returns:
tuple[int, Node] | tuple[int, int]: (0, node) if removed, (-1, None) otherwise.
Raises:
TypeError if value of unsupported type.
"""
if type(value) != self._type:
raise TypeError(f"Value must be of type {self._type}.")
status, node = self.find(value)
# If node is not in tree
if status == -1:
return -1, None
if node == self._smallest:
self._smallest = self.__find_next(node)
if node == self._largest:
self._largest = self.__find_prev(node)
# Leaf node
if node.left is None and node.right is None:
replacement = self.__delete_switch(node, None)
# 1 child (right)
elif node.left is None:
replacement = self.__delete_switch(node, node.right)
# 1 child (left)
elif node.right is None:
replacement = self.__delete_switch(node, node.left)
# 2 children
else:
# Find symmetrical successor
successor = node.right
while successor.left is not None:
successor = successor.left
if successor.parent != node:
self.__delete_switch(successor, successor.right)
successor.right = node.right
successor.right.parent = successor
self.__delete_switch(node, successor)
successor.left = node.left
successor.left.parent = successor
replacement = successor
self.__fix_remove(node, replacement)
return 0, node
[docs]
def find(self, value: object) -> tuple[int, Node | None]:
"""
Find a value in the tree.
Args:
value (object): The value to find.
Returns:
tuple[int, Node | None]: (0, current) if found, (-1, pos) otherwise.
Raises:
TypeError if value of unsupported type.
"""
if type(value) != self._type:
raise TypeError(f"Value must be of type {self._type}.")
current = self._root
pos = self._root
while current:
comp = self.__compare(value, current.data)
if not comp:
return 0, current
if comp > 0:
pos, current = current, current.right
else:
pos, current = current, current.left
return -1, pos
[docs]
def items(self) -> list[Any]:
"""
Return a list of all items in the tree ordered by their natural order.
Returns:
list[Any]: A list of all items in the tree.
"""
items: list[Any] = []
current: Node | None = self._smallest
if not current:
return items
while current != self._largest:
items.append(current.data)
current = self.__find_next(current)
items.append(self._largest.data)
return items
[docs]
def higher(self, value: object, equal_valid: bool) -> object:
"""
Find nearest successor to the value. If equal_valid is True, return the value itself if present in the tree.
Args:
value (object): The value to find the successor for.
equal_valid (bool): If True, return the value itself if present in the tree.
Returns:
object: The nearest successor to the value.
"""
found, pos = self.find(value)
if found == -1:
if self.__compare(value, pos.data) < 0:
return pos.data
if equal_valid and found == 0:
return pos.data
pos = self.__find_next(pos)
if pos:
return pos.data
return None
[docs]
def lower(self, value: object, equal_valid: bool) -> object:
"""
Find nearest predecessor to the value. If equal_valid is True, return the value itself if present in the tree.
Args:
value (object): The value to find the predecessor for.
equal_valid (bool): If True, return the value itself if present in the tree.
Returns:
object: The nearest predecessor to the value.
"""
found, pos = self.find(value)
if found == -1:
if self.__compare(value, pos.data) > 0:
return pos.data
if equal_valid and found == 0:
return pos.data
prev = self.__find_prev(pos)
if prev:
return prev.data
return None
# ----------------------------------------------------------------------
# AUXILIARY FUNCTIONS
# ----------------------------------------------------------------------
def __compare(self, val1: Any, val2: Any) -> int:
"""
Compare two values.
Args:
val1 (object): The first value to compare.
val2 (object): The second value to compare.
Returns:
int: -1 if val1 < val2, 0 if val1 == val2, 1 if val1 > val2.
"""
if val1 is None or val2 is None:
return -2
if self._key is None:
if val1 < val2:
return -1
elif val1 > val2:
return 1
else:
return 0
if self._key(val1) < self._key(val2):
return -1
elif self._key(val1) > self._key(val2):
return 1
else:
return 0
def __fix_insert(self, node: Node):
"""
Fix the tree after insertion to maintain Red-Black properties.
Args:
node (Node): The node that was just inserted.
"""
current_node = node
while current_node != self._root and current_node.parent.color == COLORS.RED:
parent_node = current_node.parent
grandparent_node = parent_node.parent
if parent_node == grandparent_node.left:
# Padre es hijo izquierdo
uncle_node = grandparent_node.right
if uncle_node and uncle_node.color == COLORS.RED:
# Recoloración (4a/4b)
parent_node.color = COLORS.BLACK
uncle_node.color = COLORS.BLACK
grandparent_node.color = COLORS.RED
current_node = grandparent_node
continue
# Necesaria rotación
if current_node == parent_node.right:
# Zigzag LR = hijo derecho (6)
current_node.color = COLORS.BLACK
grandparent_node.color = COLORS.RED
self.__rotate_lr(grandparent_node)
break
# Línea recta LL = hijo izquierdo (5)
parent_node.color = COLORS.BLACK
grandparent_node.color = COLORS.RED
self.__rotate_r(grandparent_node)
break
# Padre es hijo derecho
uncle_node = grandparent_node.left
if uncle_node and uncle_node.color == COLORS.RED:
# Recoloración (4a/4b)
parent_node.color = COLORS.BLACK
uncle_node.color = COLORS.BLACK
grandparent_node.color = COLORS.RED
current_node = grandparent_node
continue
# Necesaria rotación
if current_node == parent_node.left:
# Zigzag RL = hijo izquierdo (6)
current_node.color = COLORS.BLACK
grandparent_node.color = COLORS.RED
self.__rotate_rl(grandparent_node)
break
# Línea recta RR = hijo derecho (5)
parent_node.color = COLORS.BLACK
grandparent_node.color = COLORS.RED
self.__rotate_l(grandparent_node)
break
self._root.color = COLORS.BLACK
def __fix_remove(self, node: Node, replacement: Node | None):
"""
Fix the tree after removal to maintain Red-Black properties.
Args:
node (Node): The node that was just removed.
replacement (Node | None): The replacement node after removal.
"""
current_node = replacement
if current_node and current_node.color == COLORS.RED:
current_node.color = COLORS.BLACK
return
while current_node and current_node != self._root:
parent = current_node.parent
brother = parent.left if current_node == parent.right else parent.right
if brother and brother.color == COLORS.RED:
brother.color = COLORS.BLACK
parent.color = COLORS.RED
if brother == parent.right:
self.__rotate_l(parent)
else:
self.__rotate_r(parent)
brother = parent.right if current_node == parent.left else parent.left
if not brother:
current_node = parent
continue
closest_nephew = brother.left if parent.left == current_node else brother.right
farthest_nephew = brother.right if parent.left == current_node else brother.left
if ((not closest_nephew or closest_nephew.color == COLORS.BLACK) and
(not farthest_nephew or farthest_nephew.color == COLORS.BLACK)):
# Caso 4
brother.color = COLORS.RED
current_node = parent
continue
if ((closest_nephew and closest_nephew.color == COLORS.RED) and
(not farthest_nephew or farthest_nephew.color == COLORS.BLACK)):
# Caso 7
closest_nephew.color = COLORS.BLACK
brother.color = COLORS.RED
if brother == parent.right:
self.__rotate_r(brother)
else:
self.__rotate_l(brother)
brother = parent.right if current_node == parent.left else parent.left
if brother:
farthest_nephew = brother.right if parent.left == current_node else brother.left
if (brother and brother.color == COLORS.BLACK and
farthest_nephew and farthest_nephew.color == COLORS.RED):
# Caso 8
brother.color = parent.color
parent.color = COLORS.BLACK
farthest_nephew.color = COLORS.BLACK
if brother == parent.right:
self.__rotate_l(parent)
else:
self.__rotate_r(parent)
break
break
def __delete_switch(self, delete_node: Node, replace_node: Node | None) -> Node | None:
"""
Switch the delete node with the replace node in the tree.
Args:
delete_node (Node): The node to delete.
replace_node (Node | None): The node to replace with.
Returns:
Node | None: The replace node if it exists, otherwise None.
"""
if delete_node.parent is None:
self._root = replace_node
if replace_node is not None:
replace_node.color = COLORS.BLACK
elif self.__compare(delete_node.data, delete_node.parent.data) > 0:
delete_node.parent.right = replace_node
else:
delete_node.parent.left = replace_node
if replace_node:
replace_node.parent = delete_node.parent
delete_node.parent = None
return replace_node
def __set_parent_rotation(self, node: Node, parent: Node) -> None:
"""
Set the parent of the node after a rotation.
Args:
node (Node): The node to set the parent for.
parent (Node): The new parent of the node.
"""
node.parent = parent
if parent is None:
self._root = node
return
comp = self.__compare(node.data, parent.data)
if comp > 0:
parent.right = node
return
parent.left = node
def __rotate_r(self, node: Node) -> None:
"""
Right rotation on the given node.
Args:
node (Node): The node to rotate.
"""
subtree_parent = node.parent
pivot = node.left
node.left = pivot.right
if pivot.right:
pivot.right.parent = node
pivot.right = node
node.parent = pivot
self.__set_parent_rotation(pivot, subtree_parent)
def __rotate_l(self, node: Node) -> None:
"""
Left rotation on the given node.
Args:
node (Node): The node to rotate.
"""
subtree_parent = node.parent
pivot = node.right
node.right = pivot.left
if pivot.left:
pivot.left.parent = node
pivot.left = node
node.parent = pivot
self.__set_parent_rotation(pivot, subtree_parent)
def __rotate_lr(self, node: Node) -> None:
"""
Left-Right rotation on the given node.
Args:
node (Node): The node to rotate.
"""
pivot = node.left
child = pivot.right
if not child:
return
subtree_parent = node.parent
pivot.right = child.left
if child.left:
child.left.parent = pivot
child.left = pivot
pivot.parent = child
node.left = child.right
if child.right:
child.right.parent = node
child.right = node
node.parent = child
self.__set_parent_rotation(child, subtree_parent)
def __rotate_rl(self, node: Node) -> None:
"""
Right-Left rotation on the given node.
Args:
node (Node): The node to rotate.
"""
pivot = node.right
child = pivot.left
if not child:
return
subtree_parent = node.parent
pivot.left = child.right
if child.right:
child.right.parent = pivot
child.right = pivot
pivot.parent = child
node.right = child.left
if child.left:
child.left.parent = node
child.left = node
node.parent = child
self.__set_parent_rotation(child, subtree_parent)
def __find_next(self, node: Node):
"""
Find the successor of the given node in the tree.
Args:
node (Node): The node to find the successor for.
Returns:
Node | None: The successor node if it exists, otherwise None.
"""
if node == self._largest:
return None
if node.right:
# Sucesor en simétrico
successor = node.right
successor_parent = node
while successor:
successor_parent = successor
successor = successor.left
return successor_parent
# Sucesor para arriba
successor = node
while successor.parent and successor == successor.parent.right:
successor = successor.parent
return successor.parent
def __find_prev(self, node: Node):
"""
Find the predecessor of the given node in the tree.
Args:
node (Node): The node to find the predecessor for.
Returns:
Node | None: The predecessor node if it exists, otherwise None.
"""
if node == self._smallest:
return None
if node.left:
# Predecesor en simétrico
predecessor = node.left
predecessor_parent = node
while predecessor:
predecessor_parent = predecessor
predecessor = predecessor.right
return predecessor_parent
# Predecesor para arriba
predecessor = node
while predecessor.parent and predecessor == predecessor.parent.left:
predecessor = predecessor.parent
return predecessor.parent