Open In App

Red Black Tree in Python

Last Updated : 21 Jun, 2024
Comments
Improve
Suggest changes
Like Article
Like
Report

Red Black Tree is a self-balancing binary search tree where each node has an extra bit representing its color, either red or black. By constraining how nodes are colored during insertions and deletions, red-black trees ensure the tree remains approximately balanced during all operations, allowing for O(log⁡n) search, insert, and delete operations.

Here, we will explore, how to implement insertion, deletion, and balancing logic to construct a fully functional red-black tree in Python.

Properties of Red-Black Tree:

Red-black trees are binary search trees that add the extra constraint of node coloring to ensure the tree remains balanced during operations. Here are some key properties of red-black trees:

  1. Every node is colored red or black.
  2. The root node is always black.
  3. All leaf nodes are black and have NULL child pointers.
  4. Both children of every red node are black.
  5. Every path from a node to any of its descendant leaf nodes contains the same number of black nodes.

By constraining coloring during insertions and deletions, red-black trees ensure the maximum height of the tree does not exceed 2log⁡(n+1), where n is the number of nodes. This guarantees O(log⁡n) worst-case performance for search, insert, and delete operations.

The balance of red-black trees comes from how coloring changes during insertions and deletions. Specific balancing rules are followed that require certain coloring operations and tree rotations to retain the black-height property of the tree.

Red-Black Tree Node Implementation in Python:

We will implement a red-black tree node class that extends the typical binary search tree node structure by adding a color attribute:

  • Each node contains the standard value, left, and right attributes to represent the data value and links to children nodes.
  • The color attribute tracks if the node is ‘red’ or ‘black’. New nodes added during insertion are colored red by default.
  • The parent attribute helps track parent nodes when doing rotations and color changes.

Below is the implementation of Red-Black Tree Node in Python:

Python
class RBNode:
      # Constructor to initialize node of RB Tree
    def __init__(self, value, color='red'):
        self.value = value
        self.color = color
        self.left = None
        self.right = None
        self.parent = None

    # function to get the grandparent of node
    def grandparent(self):
        if self.parent is None:
            return None
        return self.parent.parent

    # function to get the sibling of node
    def sibling(self):
        if self.parent is None:
            return None
        if self == self.parent.left:
            return self.parent.right
        return self.parent.left

    # function to get the uncle of node
    def uncle(self):
        if self.parent is None:
            return None
        return self.parent.sibling()

3. Implementation of Red-Black Tree Insertion in Python

We will first see the algorithm for inserting new nodes into the red-black tree. This must follow specific rules to ensure the tree remains balanced. The insertion in Red Black Tree occurs in two parts:

1. Standard BST Insertion: Insertion begins by adding the new node with a red color as a standard BST insertion.

2. Fixing Red-Black Properties: After regular BST insertion, we may have violated red-black tree properties. There are two main red-black tree invariants that must be maintained after insertion:

  1. No red node can have a red child (red nodes must have black children).
  2. All paths from a node to descendant leaves must contain the same number of black nodes.

There are specific cases that can occur after adding a red child that violate these invariants:

  • Case 1: If the parent is black, the tree is still valid.
  • Case 2: If the parent is red, we have a red violation.
Python
# Class to implement Red Black Tree
class RedBlackTree:
        # Constructor
    def __init__(self):
        self.root = None

    # Function to fix the Red Black Tree properties after insertion
    def insert_fix(self, new_node):
          # While there are two continuous red nodes, we need to fix the RB tree
        while new_node.parent and new_node.parent.color == 'red':
            # If the parent is left child of grandparent
            if new_node.parent == new_node.grandparent().left:
                uncle = new_node.uncle()
                if uncle and uncle.color == 'red':
                    new_node.parent.color = 'black'
                    uncle.color = 'black'
                    new_node.grandparent().color = 'red'
                    new_node = new_node.grandparent()
                else:
                    if new_node == new_node.parent.right:
                        new_node = new_node.parent
                        self.rotate_left(new_node)
                    new_node.parent.color = 'black'
                    new_node.grandparent().color = 'red'
                    self.rotate_right(new_node.grandparent())
            # If the parent is right child of grandparent
            else:
                uncle = new_node.uncle()
                if uncle and uncle.color == 'red':
                    new_node.parent.color = 'black'
                    uncle.color = 'black'
                    new_node.grandparent().color = 'red'
                    new_node = new_node.grandparent()
                else:
                    if new_node == new_node.parent.left:
                        new_node = new_node.parent
                        self.rotate_right(new_node)
                    new_node.parent.color = 'black'
                    new_node.grandparent().color = 'red'
                    self.rotate_left(new_node.grandparent())
        self.root.color = 'black'

    # function to insert a node similar to BST insertion
    def insert(self, value):
        # Regular BST insert
        new_node = RBNode(value)
        if self.root is None:
            self.root = new_node
        else:
            curr_node = self.root
            while True:
                  # If the node is in the left subtree
                if value < curr_node.value:
                    if curr_node.left is None:
                        curr_node.left = new_node
                        new_node.parent = curr_node
                        break
                    else:
                        curr_node = curr_node.left
                # If the node is in the right subtree
                else:
                    if curr_node.right is None:
                        curr_node.right = new_node
                        new_node.parent = curr_node
                        break
                    else:
                        curr_node = curr_node.right
        self.insert_fix(new_node)

4. Implementation of Red Black Tree Rotation in Python:

The insertion fixing logic requires left and right rotations to restructure the tree in certain cases. These operations change the position of nodes to move the red violation upward while maintaining the BST ordering, similar to AVL tree rotations.

Left Rotation:

Below is the implementation of Left Rotation of RB Tree in Python:

Python
# function for left rotation of RB Tree
def rotate_left(self, node):
    right_child = node.right
    node.right = right_child.left

    if right_child.left is not None:
        right_child.left.parent = node

    right_child.parent = node.parent

    if node.parent is None:
        self.root = right_child
    elif node == node.parent.left:
        node.parent.left = right_child
    else:
        node.parent.right = right_child

    right_child.left = node
    node.parent = right_child

Right Rotation:

Below is the implementation of Right Rotation of RB Tree in Python:

Python
# function for right rotation of RB Tree
def rotate_right(self, node):
    left_child = node.left
    node.left = left_child.right

    if left_child.right is not None:
        left_child.right.parent = node

    left_child.parent = node.parent

    if node.parent is None:
        self.root = left_child
    elif node == node.parent.right:
        node.parent.right = left_child
    else:
        node.parent.left = left_child

    left_child.right = node
    node.parent = left_child

5. Implementation of Red-Black Tree Deletion in Python:

Implementing deletion in a red-black tree also requires maintaining the coloring invariants after removing nodes. Deletion in Red Black Tree occurs in two parts:

1. Standard BST Deletion: We first implement the standard BST node removal process.

2. Red-Black Tree Fixing After Deletion: Similar to insertion, a deletion can break the red-black tree invariants. There are two main cases:

  • Case 1: Deleted node was black.
  • Case 2: Deleted node was red.

Below is the implementation of Deletion in Red Black Tree in Python:

Python
def delete(self, value):
    node_to_remove = self.search(value)

    if node_to_remove is None:
        return

    if node_to_remove.left is None or node_to_remove.right is None:
        self._replace_node(
            node_to_remove, node_to_remove.left or node_to_remove.right)
    else:
        successor = self._find_min(node_to_remove.right)
        node_to_remove.value = successor.value
        self._replace_node(successor, successor.right)

    self.delete_fix(node_to_remove)

def delete_fix(self, x):
    while x != self.root and x.color == 'black':
        if x == x.parent.left:
            sibling = x.sibling()
            if sibling.color == 'red':
                sibling.color = 'black'
                x.parent.color = 'red'
                self.rotate_left(x.parent)
                sibling = x.sibling()
            if (sibling.left is None or sibling.left.color == 'black') and (sibling.right is None or sibling.right.color == 'black'):
                sibling.color = 'red'
                x = x.parent
            else:
                if sibling.right is None or sibling.right.color == 'black':
                    sibling.left.color = 'black'
                    sibling.color = 'red'
                    self.rotate_right(sibling)
                    sibling = x.sibling()
                sibling.color = x.parent.color
                x.parent.color = 'black'
                if sibling.right:
                    sibling.right.color = 'black'
                self.rotate_left(x.parent)
                x = self.root
        else:
            sibling = x.sibling()
            if sibling.color == 'red':
                sibling.color = 'black'
                x.parent.color = 'red'
                self.rotate_right(x.parent)
                sibling = x.sibling()
            if (sibling.left is None or sibling.left.color == 'black') and (sibling.right is None or sibling.right.color == 'black'):
                sibling.color = 'red'
                x = x.parent
            else:
                if sibling.left is None or sibling.left.color == 'black':
                    sibling.right.color = 'black'
                    sibling.color = 'red'
                    self.rotate_left(sibling)
                    sibling = x.sibling()
                sibling.color = x.parent.color
                x.parent.color = 'black'
                if sibling.left:
                    sibling.left.color = 'black'
                self.rotate_right(x.parent)
                x = self.root
    x.color = 'black'

Complete Red-Black Tree Implementation in Python

Putting everything together, here is an implementation of a complete functional red-black tree in Python:

Python
# class to implement node of RB Tree
class RBNode:
        # cnostructor
    def __init__(self, value, color='red'):
        self.value = value
        self.color = color
        self.left = None
        self.right = None
        self.parent = None

    # function to get the grandparent of node
    def grandparent(self):
        if self.parent is None:
            return None
        return self.parent.parent

    # function to get the sibling of node
    def sibling(self):
        if self.parent is None:
            return None
        if self == self.parent.left:
            return self.parent.right
        return self.parent.left

    # function to get the uncle of node
    def uncle(self):
        if self.parent is None:
            return None
        return self.parent.sibling()

# function to implement Red Black Tree


class RedBlackTree:
        # constructor to initialize the RB tree
    def __init__(self):
        self.root = None

    # function to search a value in RB Tree
    def search(self, value):
        curr_node = self.root
        while curr_node is not None:
            if value == curr_node.value:
                return curr_node
            elif value < curr_node.value:
                curr_node = curr_node.left
            else:
                curr_node = curr_node.right
        return None

    # function to insert a node in RB Tree, similar to BST insertion
    def insert(self, value):
        # Regular insertion
        new_node = RBNode(value)
        if self.root is None:
            self.root = new_node
        else:
            curr_node = self.root
            while True:
                if value < curr_node.value:
                    if curr_node.left is None:
                        curr_node.left = new_node
                        new_node.parent = curr_node
                        break
                    else:
                        curr_node = curr_node.left
                else:
                    if curr_node.right is None:
                        curr_node.right = new_node
                        new_node.parent = curr_node
                        break
                    else:
                        curr_node = curr_node.right
        self.insert_fix(new_node)

    # Function to fix RB tree properties after insertion
    def insert_fix(self, new_node):
        while new_node.parent and new_node.parent.color == 'red':
            if new_node.parent == new_node.grandparent().left:
                uncle = new_node.uncle()
                if uncle and uncle.color == 'red':
                    new_node.parent.color = 'black'
                    uncle.color = 'black'
                    new_node.grandparent().color = 'red'
                    new_node = new_node.grandparent()
                else:
                    if new_node == new_node.parent.right:
                        new_node = new_node.parent
                        self.rotate_left(new_node)
                    new_node.parent.color = 'black'
                    new_node.grandparent().color = 'red'
                    self.rotate_right(new_node.grandparent())
            else:
                uncle = new_node.uncle()
                if uncle and uncle.color == 'red':
                    new_node.parent.color = 'black'
                    uncle.color = 'black'
                    new_node.grandparent().color = 'red'
                    new_node = new_node.grandparent()
                else:
                    if new_node == new_node.parent.left:
                        new_node = new_node.parent
                        self.rotate_right(new_node)
                    new_node.parent.color = 'black'
                    new_node.grandparent().color = 'red'
                    self.rotate_left(new_node.grandparent())
        self.root.color = 'black'

    # function to delete a value from RB Tree
    def delete(self, value):
        node_to_remove = self.search(value)

        if node_to_remove is None:
            return

        if node_to_remove.left is None or node_to_remove.right is None:
            self._replace_node(
                node_to_remove, node_to_remove.left or node_to_remove.right)
        else:
            successor = self._find_min(node_to_remove.right)
            node_to_remove.value = successor.value
            self._replace_node(successor, successor.right)

        self.delete_fix(node_to_remove)

    # function to fix RB Tree properties after deletion
    def delete_fix(self, x):
        while x != self.root and x.color == 'black':
            if x == x.parent.left:
                sibling = x.sibling()
                if sibling.color == 'red':
                    sibling.color = 'black'
                    x.parent.color = 'red'
                    self.rotate_left(x.parent)
                    sibling = x.sibling()
                if (sibling.left is None or sibling.left.color == 'black') and (sibling.right is None or sibling.right.color == 'black'):
                    sibling.color = 'red'
                    x = x.parent
                else:
                    if sibling.right is None or sibling.right.color == 'black':
                        sibling.left.color = 'black'
                        sibling.color = 'red'
                        self.rotate_right(sibling)
                        sibling = x.sibling()
                    sibling.color = x.parent.color
                    x.parent.color = 'black'
                    if sibling.right:
                        sibling.right.color = 'black'
                    self.rotate_left(x.parent)
                    x = self.root
            else:
                sibling = x.sibling()
                if sibling.color == 'red':
                    sibling.color = 'black'
                    x.parent.color = 'red'
                    self.rotate_right(x.parent)
                    sibling = x.sibling()
                if (sibling.left is None or sibling.left.color == 'black') and (sibling.right is None or sibling.right.color == 'black'):
                    sibling.color = 'red'
                    x = x.parent
                else:
                    if sibling.left is None or sibling.left.color == 'black':
                        sibling.right.color = 'black'
                        sibling.color = 'red'
                        self.rotate_left(sibling)
                        sibling = x.sibling()
                    sibling.color = x.parent.color
                    x.parent.color = 'black'
                    if sibling.left:
                        sibling.left.color = 'black'
                    self.rotate_right(x.parent)
                    x = self.root
        x.color = 'black'

    # Function for left rotation of RB Tree
    def rotate_left(self, node):
        right_child = node.right
        node.right = right_child.left

        if right_child.left is not None:
            right_child.left.parent = node

        right_child.parent = node.parent

        if node.parent is None:
            self.root = right_child
        elif node == node.parent.left:
            node.parent.left = right_child
        else:
            node.parent.right = right_child

        right_child.left = node
        node.parent = right_child

    # function for right rotation of RB Tree
    def rotate_right(self, node):
        left_child = node.left
        node.left = left_child.right

        if left_child.right is not None:
            left_child.right.parent = node

        left_child.parent = node.parent

        if node.parent is None:
            self.root = left_child
        elif node == node.parent.right:
            node.parent.right = left_child
        else:
            node.parent.left = left_child

        left_child.right = node
        node.parent = left_child

    # function to replace an old node with a new node
    def _replace_node(self, old_node, new_node):
        if old_node.parent is None:
            self.root = new_node
        else:
            if old_node == old_node.parent.left:
                old_node.parent.left = new_node
            else:
                old_node.parent.right = new_node
        if new_node is not None:
            new_node.parent = old_node.parent

    # function to find node with minimum value in a subtree
    def _find_min(self, node):
        while node.left is not None:
            node = node.left
        return node

    # function to perform inorder traversal
    def _inorder_traversal(self, node):
        if node is not None:
            self._inorder_traversal(node.left)
            print(node.value, end=" ")
            self._inorder_traversal(node.right)


# Example driver code
if __name__ == "__main__":
    tree = RedBlackTree()
    tree.insert(10)
    tree.insert(20)
    tree.insert(30)
    tree.insert(40)
    tree.insert(50)
    tree.insert(25)

    print("Inorder traversal of the Red-Black Tree:")
    tree._inorder_traversal(tree.root)
    print()

    tree.delete(20)

    print("Inorder traversal of the Red-Black Tree after deleting 20")
    tree._inorder_traversal(tree.root)
    print()

Output
Inorder traversal of the Red-Black Tree:
10 20 25 30 40 50 
Inorder traversal of the Red-Black Tree after deleting 20
10 25 30 40 50 



Next Article
Article Tags :
Practice Tags :

Similar Reads