Red-Black Tree Implementation in Python
A red-black tree is a self-balancing binary search tree where each node stores an extra bit representing “color” (“red” or “black”). The properties of red-black trees ensure that the tree remains approximately balanced, making the insertion, deletion, and search operations efficient with O(log n) time complexity.
Program Structure
The implementation of the red-black tree consists of the following components:
1. Node Class
This class represents a node in the red-black tree. It includes the following attributes:
key: The key stored in the node.color: The color of the node (either “red” or “black”).left,right,parent: Pointers to the left child, right child, and parent nodes, respectively.
2. RedBlackTree Class
This class encapsulates the functionality of the red-black tree and includes the following methods:
__init__(self): Initializes the red-black tree, including the root and the sentinel node.insert(self, key): Inserts a new key into the red-black tree and balances the tree according to red-black properties.insert_fixup(self, z): Fixes any violations of red-black properties after an insertion.left_rotate(self, x): Performs a left rotation around a given node.right_rotate(self, x): Performs a right rotation around a given node.search(self, key): Searches for a key in the red-black tree.
Python Code Implementation
class Node:
def __init__(self, key, color="red"):
"""
Initialize a new node.
:param key: The key value of the node.
:param color: The color of the node (red or black).
"""
self.key = key
self.color = color # Node color, default to red.
self.left = None
self.right = None
self.parent = None
class RedBlackTree:
def __init__(self):
"""
Initialize the red-black tree.
"""
self.NIL = Node(key=None, color="black") # Sentinel node represents leaves.
self.root = self.NIL
def left_rotate(self, x):
"""
Perform a left rotation around the node x.
:param x: The node to rotate around.
"""
y = x.right
x.right = y.left
if y.left != self.NIL:
y.left.parent = x
y.parent = x.parent
if x.parent is None:
self.root = y
elif x == x.parent.left:
x.parent.left = y
else:
x.parent.right = y
y.left = x
x.parent = y
def right_rotate(self, x):
"""
Perform a right rotation around the node x.
:param x: The node to rotate around.
"""
y = x.left
x.left = y.right
if y.right != self.NIL:
y.right.parent = x
y.parent = x.parent
if x.parent is None:
self.root = y
elif x == x.parent.right:
x.parent.right = y
else:
x.parent.left = y
y.right = x
x.parent = y
def insert(self, key):
"""
Insert a new key into the red-black tree.
:param key: The key to insert.
"""
z = Node(key)
z.left = self.NIL
z.right = self.NIL
y = None
x = self.root
while x != self.NIL:
y = x
if z.key < x.key:
x = x.left
else:
x = x.right
z.parent = y
if y is None:
self.root = z
elif z.key < y.key:
y.left = z
else:
y.right = z
z.color = "red"
self.insert_fixup(z)
def insert_fixup(self, z):
"""
Fix the red-black tree after insertion to maintain properties.
:param z: The node that was inserted.
"""
while z.parent and z.parent.color == "red":
if z.parent == z.parent.parent.left:
y = z.parent.parent.right
if y.color == "red": # Case 1
z.parent.color = "black"
y.color = "black"
z.parent.parent.color = "red"
z = z.parent.parent
else:
if z == z.parent.right: # Case 2
z = z.parent
self.left_rotate(z)
z.parent.color = "black" # Case 3
z.parent.parent.color = "red"
self.right_rotate(z.parent.parent)
else:
y = z.parent.parent.left
if y.color == "red": # Case 1
z.parent.color = "black"
y.color = "black"
z.parent.parent.color = "red"
z = z.parent.parent
else:
if z == z.parent.left: # Case 2
z = z.parent
self.right_rotate(z)
z.parent.color = "black" # Case 3
z.parent.parent.color = "red"
self.left_rotate(z.parent.parent)
self.root.color = "black"
def search(self, key, node=None):
"""
Search for a key in the red-black tree.
:param key: The key to search for.
:param node: The node to start the search from (defaults to root).
:return: The node containing the key, or None if not found.
"""
if node is None:
node = self.root
while node != self.NIL and key != node.key:
if key < node.key:
node = node.left
else:
node = node.right
return node if node != self.NIL else None
Explanation
The red-black tree implementation is centered around two main classes: Node and RedBlackTree.
1. Node Class
The Node class represents a node in the red-black tree, containing:
key: The key stored in the node.color: The color of the node, either “red” or “black”.left,right,parent: Pointers to the left child, right child, and parent nodes, respectively.
2. RedBlackTree Class
The RedBlackTree class manages the red-black tree and provides the following key functionalities:
Insert Method
The insert method adds a new key to the tree. It starts by placing the new node as a red leaf and then calls insert_fixup to ensure that the red-black properties are maintained.
Insert Fixup
The insert_fixup method is crucial for maintaining the balance of the tree. It handles the three cases that arise during insertion:
- Case 1: The new node’s parent and its uncle are both red. The solution is to recolor them and move the problem up the tree.
- Case 2: The new node’s parent is red, but the uncle is black, and the new node is on the opposite side from its parent. A rotation is needed to transform the situation into Case 3.
- Case 3: The new node’s parent is red, but the uncle is black, and the new node is on the same side as its parent. A rotation and recoloring fix the violation.
Rotations
The left_rotate and right_rotate methods perform rotations to adjust the tree structure. Rotations are necessary to maintain the balanced nature of the red-black tree during insertion and deletion operations.
Search Method
The search method allows searching for a key in the tree. It follows the binary search tree property to locate the key, returning the node if found, or None if the key is not present.
Usage Example
# Example usage of the RedBlackTree
# Create a red-black tree
rbt = RedBlackTree()
# Insert some keys into the tree
keys = [20, 15, 25, 10, 5, 1]
for key in keys:
rbt.insert(key)
# Search for a key in the tree
result = rbt.search(15)
if result:
print(f"Key 15 found with color: {result.color}")
else:
print("Key 15 not found in the tree.")
# Output: Key 15 found with color: black
This example demonstrates how to create a red-black tree, insert keys into it, and search for a specific key. The output shows whether the key is found and its color, indicating the balanced nature of the tree.
