Implement Red-Black Tree in Python - Guide for Technical Coding Interview

A red-black tree is a self-balancing binary search tree where each node has an extra bit representing 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.

Red-black trees are important data structures widely used in languages like Java and C++. Here we will examine how to implement insertion, deletion, and balancing logic to construct a fully functional red-black tree in Python.

Open Table of Contents

Overview of Red-Black Trees
Red-Black Tree Node Implementation
Implementing Red-Black Tree Insertion
Deletion in a Red-Black Tree
Complete Red-Black Tree in Python
Applications of Red-Black Trees

Overview of Red-Black Trees

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

Every node is colored red or black.
The root node is always black.
All leaf nodes are black with NULL child pointers.
Both children of every red node are black.
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 (number of black nodes on any root-to-leaf path) 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

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

class RBNode:
  def __init__(self, value, color='red'):
    self.value = value
    self.color = color
    self.left = None
    self.right = None
    self.parent = None

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 is helpful for tracking parent nodes when doing rotations and color changes.

Implementing Red-Black Tree Insertion

We will first examine the algorithm for inserting new nodes into the red-black tree. This must follow specific rules to ensure the tree remains balanced.

The steps for insertion are:

Perform standard BST insertion, adding the node as a red leaf node.
Fix red-black properties if needed through color changes and rotations.

Standard BST Insertion

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

def insert(self, value):
  # Regular BST insert
  new_node = RBNode(value)
  if self.root is None:
    # Handle empty tree
    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

This follows the standard recursive BST insertion logic, descending the tree and adding the new node with the given value once we reach a leaf node.

The new node is initially added with a red color.

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:

No red node can have a red child (so red nodes must have black children).
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, tree is still valid.

If the parent of the new red node is black, both properties are satisfied. Nothing needs to be done since rules are not violated.

Case 2: If parent is red, we have a red violation.

If the parent is red, this means there is now a red node with a red child, violating property 1. To fix this, there are subcases:

Case 2a: If parent’s sibling is red - Recolor the parent and sibling to black, and grandparent to red. This maintains a valid tree, but grandparent may now be red, so recursively check it.
Case 2b: Parent’s sibling is black - Perform rotations to restructure the tree so the red violation moves up to fix.

The full fixing logic including handling these cases is shown below:

def insert_fix(self, new_node):
  # Fix red-black properties if needed
  parent = new_node.parent

  # Case 1: If parent is black, tree is still valid
  if parent is None or parent.color == 'black':
    return

  # Case 2: If parent is red, have to handle violations
  else:
    if new_node.uncle() is not None and new_node.uncle().color == 'red':
      # Case 2a: Uncle is red, recolor parent, grandparent, uncle
      new_node.parent.color = 'black'
      new_node.uncle().color = 'black'
      new_node.grandparent().color = 'red'

      # Recurse up the tree
      self.insert_fix(new_node.grandparent())

    else:
      # Case 2b: Uncle is black (or null), restructure tree
      grandparent = new_node.grandparent()

      # New node is left child of parent, parent is left of grandparent
      if new_node == new_node.parent.left:
        if new_node.parent == grandparent.left:
          # Left-left case
          self.rotate_right(grandparent)
          new_node.parent.color = 'black'
          grandparent.color = 'red'

        else:
          # Left-right case
          self.rotate_left(new_node.parent)
          self.rotate_right(grandparent)
          new_node.color = 'black'
          grandparent.color = 'red'

      # Mirror cases: parent/new node on right
      else:
        if new_node.parent == grandparent.right:
          # Right-right case
          self.rotate_left(grandparent)
          new_node.parent.color = 'black'
          grandparent.color = 'red'

        else:
          # Right-left case
          self.rotate_right(new_node.parent)
          self.rotate_left(grandparent)
          new_node.color = 'black'
          grandparent.color = 'red'

This covers all cases that can occur after adding a red child node, ensuring red-black properties are restored through rotations and recoloring.

The full insert method handles both standard insertion and fixing:

def insert(self, value):
  # Regular BST insert
  new_node = RBNode(value)
  # Code for standard insertion...

  # Fix red-black properties
  self.insert_fix(new_node)

Rotations for Balancing

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.

Here is an implementation of left and right rotations on a given node:

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:
    # Node is root
    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


def rotate_right(self, node):
  # Mirror of rotate left
  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

These rotations restructure the nodes such that the rotated node becomes a child while its child becomes the new parent. This moves a red violation upward in the tree.

Deletion in a Red-Black Tree

Implementing deletion in a red-black tree also requires maintaining the coloring invariants after removing nodes.

The steps for deletion are:

Perform standard BST deletion logic of removing the node and restructuring children.
Fix any red-black violations through rotations and recoloring.

Standard BST Deletion

We first implement the standard BST node removal process:

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

  if node_to_remove is None:
    # Value not found, no deletion needed
    return

  if node_to_remove.left is None:
    self._replace_node(node_to_remove, node_to_remove.right)
  elif node_to_remove.right is None:
    self._replace_node(node_to_remove, node_to_remove.left)
  else:
    # Node has two children
    # Find successor (leftmost of right subtree)
    successor = self._find_min(node_to_remove.right)
    successor_value = successor.value
    self._replace_node(successor, successor.right)
    node_to_remove.value = successor_value

  # Fix red-black properties if needed
  self.delete_fix(node_to_remove)

This follows the same logic as a standard BST deletion. If the node is a leaf, it can simply be removed. If it has one child, that child replaces it. If it has two children, we find its successor in the right subtree and swap their values before deleting the successor leaf node instead.

After removal, we check if red-black properties were violated and call fix-up logic.

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

If the deleted node was black, removing it caused black-heights of paths to diverge, since other paths did not have their black node count reduced.

To balance this, we check the color of the deleted node’s child. If it is red, recolor it black to maintain black heights. If it is black, more complex fixes are needed:

Case 1a: Red sibling - Rotate to move the red up to parent level. This converts it to one of the other cases.
Case 1b: Black sibling with black children - Recolor sibling red and propagate the double black up the tree.
Case 1c: Black sibling with at least 1 red child - Recolor to balance black counts and preserve properties.

Case 2: Deleted node was red

If the deleted node was red, since red nodes do not add to the black height of paths we know heights remain valid.

Here is the full deletion fixup logic:

def delete_fix(self, x):

  # Start fixes from parent since x is gone
  x_parent = x.parent

  if x.color == 'red':
    # Case 2: If x was red, tree is still valid
    return

  if x_parent is None:
    # Case 1: x was black, but was root
    # Set new root to black
    self.root.color = 'black'
    return

  # Case 1: x was black
  if x.sibling().color == 'red':
    # Case 1a: Sibling is red
    x_parent.color = 'red'
    x.sibling().color = 'black'
    if x == x_parent.left:
      self.rotate_left(x_parent)
    else:
      self.rotate_right(x_parent)

  # Fix double black at x_parent
  self.delete_fix(x_parent)

def delete_fix(self, x):

  if x.color == 'red':
    # Color flip
    x.color = 'black'
    return

  if x.parent is None:
    # Reached root, all black-heights fixed
    return

  if x.sibling().color == 'black':
    # Case 1b and 1c
    if (x.sibling().left is not None and x.sibling().left.color == 'red') or (x.sibling().right is not None and x.sibling().right.color == 'red'):
      # Case 1c: At least 1 red child
      if x.sibling().left is not None and x.sibling().left.color == 'red':
        if x == x.parent.left:
          self.rotate_right(x.sibling())
          x.sibling().color = x.parent.color
          x.parent.color = 'black'
        else:
          x.sibling().left.color = x.parent.color
          x.parent.color = 'black'
          self.rotate_right(x.sibling())
          self.rotate_left(x.parent)
      else:
        if x == x.parent.left:
          x.sibling().right.color = x.parent.color
          x.parent.color = 'black'
          self.rotate_left(x.sibling())
          self.rotate_right(x.parent)
        else:
          self.rotate_left(x.sibling())
          x.sibling().color = x.parent.color
          x.parent.color = 'black'
    else:
      # Case 1b: All black children
      x.sibling().color = 'red'
      if x.parent.color == 'black':
        self.delete_fix(x.parent)
      else:
        x.parent.color = 'black'

  # Pass double black node up tree
  self.delete_fix(x.parent)

This comprehensive fixing logic handles all cases after deletion to restore red-black properties. The tree remains balanced after any insertion or deletion.

Full Deletion Method

Combining the standard deletion and fixing logic, the full delete method is:

def delete(self, value):

  # Standard deletion
  node_to_remove = self.search(value)
  # Delete node cases

  # Fix red-black properties
  self.delete_fix(node_to_remove)

Complete Red-Black Tree in Python

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

import sys

class RBNode:
  # Red-Black Tree Node class

  def __init__(self, value, color='red'):
    ...

  # Required helper methods and properties
  def grandparent(self): ...

  def sibling(self): ...

  def uncle(self): ...


class RedBlackTree:

  def __init__(self):
    self.root = None

  def search(self, value):
    # Search logic

  def insert(self, value):
    # Regular insertion
    new_node = RBNode(value)
    # Insert fixing
    self.insert_fix(new_node)

  def insert_fix(self, new_node):
    # Full insert fix logic

  def delete(self, value):
    # Standard deletion
    node_to_remove = self.search(value)
    # Delete node cases

    # Fix red-black properties
    self.delete_fix(node_to_remove)

  def delete_fix(self, x):
    # Full delete fix logic

  # Rotation methods
  def rotate_left(self, node):
    ...

  def rotate_right(self, node):
    ...

This provides a full implementation of a red-black tree data structure in Python. The key additions to a standard BST are the color tracking, the insertion fix and deletion fix logic, and the left/right rotations used to rebalance during insertions and deletions.

Applications of Red-Black Trees

Red-black trees are used extensively in various systems and libraries where balanced search trees are required. Some example applications include:

Implementing Map/Set data structures - In languages like C++ and Java, the std::map and HashSet classes use red-black trees internally for sorting keys and providing O(log n) operations.
Compiler symbol tables - To store identifiers in programming languages for quick lookup time during compilation.
Networking and routing algorithms - To store routing table info while allowing fast route lookups and updates.
Scheduling algorithms - The Linux Completely Fair Scheduler uses red-black trees to store dynamic process info sorted by priority.
Floating point arithmetic - The binary exponent can be stored in a red-black tree for fast floating point math operations.
Scientific computing - Fast O(log n) ordered set operations provided by red-black trees are useful for problems like computational geometry and mesh generation.

For these applications, red-black trees provide optimal O(log n) performance for search, insert, and delete while maintaining balanced tree height through their unique coloring algorithm. This makes them one of the most effective balanced binary search tree structures.