Building a Red-Black Binary Tree in Python
Tags: python, compsci, algorithms, binarytree
Published: 2021-06-28 14:52:16

---
title: Building a Red-Black Binary Tree in Python
published: true
date: 2021-06-21 13:02:49 UTC
tags: python,compsci,algorithms,binarytree
canonical_url: https://qvault.io/python/red-black-tree-python/
---

![red black](https://qvault.io/wp-content/uploads/2021/06/red-black.jpg "red black")

The post [Building a Red-Black Binary Tree in Python](https://qvault.io/python/red-black-tree-python/) first appeared on [Qvault](https://qvault.io).

A red-black tree is a kind of self-balancing binary search tree. Each node stores an extra bit, which we will call the color, red or black. The color ensures that the tree remains _approximately_ balanced during insertions and deletions. When the tree is modified, the new tree is rearranged and repainted to restore the coloring properties that constrain how unbalanced the tree can become in the worst case.

The purpose of a red-black tree is to stay balanced which ensures that its common operations, like lookup and delete, never degrade to worse than `O(n*log(n))`.

Sorry to interrupt! I just wanted to mention that you should check out my new free Python course, “Big-O Data Structures”. You’ll learn all you need to know to crush your upcoming coding interviews and whiteboard sessions.

[Start Python Data Structures Course](https://qvault.io/big-o-data-structures-course/)

## What is a balanced binary tree?

Since the reason colors are added to a binary tree is to ensure that it remains balanced, we need to understand how and why a binary tree is balanced. To put it simply, a balanced tree’s branches differ in height by no more than 1.

The following tree is balanced because between its two branches one has a height of 2, and the other 3, meaning they differ by no more than 1.

```
A
   / \
  B C 
 /
D
```

The next tree is unbalanced because it’s branches differ in height by more than 1. `C`‘s right side has a height of 2 while its left side has a height of 4).

```
A
   / \
  B C
 / /
D E  
     /  
    G
```

## Why do we want balanced trees?

Balanced binary search trees ensure **speed**. The speed of an operation in a binary tree depends on the height of the tree. If the tree is balanced, then the height is only the `log` of the number of nodes, which means the tree will work as fast as possible. However, if the tree is unbalanced, for example with one really long branch, then the height because the total number of nodes rather than the log.

```
A
     / 
    B
   /
  C
 /  
D
```

## Properties of a red-black tree

In addition to all the properties of a [Binary Search Tree](https://qvault.io/python/binary-search-tree-in-python/), a red-black tree must have the following:

1. Each node is either red or black
2. The root is black. This rule is sometimes omitted. Since the root can always be changed from red to black, but not necessarily vice versa, this rule has little effect on analysis.
3. All `nil` leaf nodes are black.
4. If a node is red, then both its children are black.
5. All paths from a single node go through the same number of black nodes in order to reach any of its descendant `nil` nodes.

## Implementing a Red-Black Tree in Python

### Step 1 – RBNode Class

Our implementation will use a `Tree` class and a `Node` class. The Node will be fairly simple, it’s just a constructor.

```python
class RBNode:
    def __init__ (self, val):
        self.red = False
        self.parent = None
        self.val = val
        self.left = None
        self.right = None
<small id="shcb-language-1"><span>Code language:</span> <span>Python</span> <span>(</span><span>python</span><span>)</span></small>
```

### Step 2 – RBTree Class

Next let’s create a tree class with a constructor.

```
class RBTree:
    def __init__ (self):
        self.nil = RBNode(0)
        self.nil.red = False
        self.nil.left = None
        self.nil.right = None
        self.root = self.nil`
<small id="shcb-language-2"><span>Code language:</span> <span>Python</span> <span>(</span><span>python</span><span>)</span></small>
```

### Step 3 – Insert method

```python
def insert(self, val):
    # Ordinary Binary Search Insertion
    new_node = RBNode(val)
    new_node.parent = None
    new_node.left = self.nil
    new_node.right = self.nil
    new_node.red = True # new node must be red

    parent = None
    current = self.root
    while current != self.nil:
        parent = current
        if new_node.val < current.val:
            current = current.left
        elif new_node.val > current.val:
            current = current.right
        else:
            return

    # Set the parent and insert the new node
    new_node.parent = parent
    if parent == None:
        self.root = new_node
    elif new_node.val < parent.val:
        parent.left = new_node
    else:
        parent.right = new_node

    # Fix the tree
    self.fix_insert(new_node)
<small id="shcb-language-3"><span>Code language:</span> <span>Python</span> <span>(</span><span>python</span><span>)</span></small>
```

The insert method will look a lot like a traditional binary tree insert method. The biggest difference is that after doing an insert, we’ll call a special `fix_insert` method. For now just call it, we’ll implement it in just a moment.

### Step 4 – Rotate left

We’ll need some rotation methods in our “fix” step that’s coming up. Let’s code those now.

![rotate red black tree right](https://qvault.io/wp-content/uploads/2021/06/rotate_red_black_tree_right.gif "rotate red black tree right")

```python
# rotate left at node x
def rotate_left(self, x):
    y = x.right
    x.right = y.left
    if y.left != self.nil:
        y.left.parent = x

    y.parent = x.parent
    if x.parent == None:
        self.root = y
    elif x == x.parent.left:
        x.parent.left = y
    else:
        x.parent.right = y
    y.left = x
    x.parent = y
<small id="shcb-language-4"><span>Code language:</span> <span>Python</span> <span>(</span><span>python</span><span>)</span></small>
```

### Step 5 – Rotate right

```python
# rotate right at node x
def rotate_right(self, x):
    y = x.left
    x.left = y.right
    if y.right != self.nil:
        y.right.parent = x

    y.parent = x.parent
    if x.parent == None:
        self.root = y
    elif x == x.parent.right:
        x.parent.right = y
    else:
        x.parent.left = y
    y.right = x
    x.parent = y
<small id="shcb-language-5"><span>Code language:</span> <span>Python</span> <span>(</span><span>python</span><span>)</span></small>
```

## Step 6 – Fix insert

The real bread and butter is in this step, it’s what makes a red-black tree balanced.

```python
def fix_insert(self, new_node):
    while new_node != self.root and new_node.parent.red:
        if new_node.parent == new_node.parent.parent.right:
            u = new_node.parent.parent.left # uncle
            if u.red:

                u.red = False
                new_node.parent.red = False
                new_node.parent.parent.red = True
                new_node = new_node.parent.parent
            else:
                if new_node == new_node.parent.left:
                    new_node = new_node.parent
                    self.rotate_right(new_node)
                new_node.parent.red = False
                new_node.parent.parent.red = True
                self.rotate_left(new_node.parent.parent)
        else:
            u = new_node.parent.parent.right # uncle

            if u.red:
                u.red = False
                new_node.parent.red = False
                new_node.parent.parent.red = True
                new_node = new_node.parent.parent
            else:
                if new_node == new_node.parent.right:
                    new_node = new_node.parent
                    self.rotate_left(new_node)
                new_node.parent.red = False
                new_node.parent.parent.red = True
                self.rotate_right(new_node.parent.parent)
    self.root.red = False
<small id="shcb-language-6"><span>Code language:</span> <span>Python</span> <span>(</span><span>python</span><span>)</span></small>
```

## Full Example of Red-Black Tree in Code

```python
import random

class RBNode:
    def __init__ (self, val):
        self.red = False
        self.parent = None
        self.val = val
        self.left = None
        self.right = None

class RBTree:
    def __init__ (self):
        self.nil = RBNode(0)
        self.nil.red = False
        self.nil.left = None
        self.nil.right = None
        self.root = self.nil

    def insert(self, val):
        # Ordinary Binary Search Insertion
        new_node = RBNode(val)
        new_node.parent = None
        new_node.left = self.nil
        new_node.right = self.nil
        new_node.red = True # new node must be red

        parent = None
        current = self.root
        while current != self.nil:
            parent = current
            if new_node.val < current.val:
                current = current.left
            elif new_node.val > current.val:
                current = current.right
            else:
                return

        # Set the parent and insert the new node
        new_node.parent = parent
        if parent == None:
            self.root = new_node
        elif new_node.val < parent.val:
            parent.left = new_node
        else:
            parent.right = new_node

        # Fix the tree
        self.fix_insert(new_node)

    def fix_insert(self, new_node):
        while new_node != self.root and new_node.parent.red:
            if new_node.parent == new_node.parent.parent.right:
                u = new_node.parent.parent.left # uncle
                if u.red:
                    u.red = False
                    new_node.parent.red = False
                    new_node.parent.parent.red = True
                    new_node = new_node.parent.parent
                else:
                    if new_node == new_node.parent.left:
                        new_node = new_node.parent
                        self.rotate_right(new_node)
                    new_node.parent.red = False
                    new_node.parent.parent.red = True
                    self.rotate_left(new_node.parent.parent)
            else:
                u = new_node.parent.parent.right # uncle

                if u.red:
                    u.red = False
                    new_node.parent.red = False
                    new_node.parent.parent.red = True
                    new_node = new_node.parent.parent
                else:
                    if new_node == new_node.parent.right:
                        new_node = new_node.parent
                        self.rotate_left(new_node)
                    new_node.parent.red = False
                    new_node.parent.parent.red = True
                    self.rotate_right(new_node.parent.parent)
        self.root.red = False

    def exists(self, val):
        curr = self.root
        while curr != self.nil and val != curr.val:
            if val < curr.val:
                curr = curr.left
            else:
                curr = curr.right
        return curr

    # rotate left at node x
    def rotate_left(self, x):
        y = x.right
        x.right = y.left
        if y.left != self.nil:
            y.left.parent = x

        y.parent = x.parent
        if x.parent == None:
            self.root = y
        elif x == x.parent.left:
            x.parent.left = y
        else:
            x.parent.right = y
        y.left = x
        x.parent = y

    # rotate right at node x
    def rotate_right(self, x):
        y = x.left
        x.left = y.right
        if y.right != self.nil:
            y.right.parent = x

        y.parent = x.parent
        if x.parent == 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 __repr__ (self):
        lines = []
        print_tree(self.root, lines)
        return '\n'.join(lines)

def print_tree(node, lines, level=0):
    if node.val != 0:
        print_tree(node.left, lines, level + 1)
        lines.append('-' * 4 * level + '> ' +
                     str(node.val) + ' ' + ('r' if node.red else 'b'))
        print_tree(node.right, lines, level + 1)

def get_nums(num):
    random.seed(1)
    nums = []
    for _ in range(num):
        nums.append(random.randint(1, num-1))
    return nums

def main():
    tree = RBTree()
    for x in range(1, 51):
        tree.insert(x)
    print(tree)

main()
<small id="shcb-language-7"><span>Code language:</span> <span>Python</span> <span>(</span><span>python</span><span>)</span></small>
```

Ready to get coding?

[<svg aria-hidden="true" role="img" height="1em" width="1em" viewbox="0 0 256 512" xmlns="http://www.w3.org/2000/svg"><path fill="currentColor" d="M224.3 273l-136 136c-9.4 9.4-24.6 9.4-33.9 0l-22.6-22.6c-9.4-9.4-9.4-24.6 0-33.9l96.4-96.4-96.4-96.4c-9.4-9.4-9.4-24.6 0-33.9L54.3 103c9.4-9.4 24.6-9.4 33.9 0l136 136c9.5 9.4 9.5 24.6.1 34z"></path></svg>Try our coding courses free](https://app.qvault.io/)

[<svg aria-hidden="true" role="img" height="1em" width="1em" viewbox="0 0 256 512" xmlns="http://www.w3.org/2000/svg"><path fill="currentColor" d="M224.3 273l-136 136c-9.4 9.4-24.6 9.4-33.9 0l-22.6-22.6c-9.4-9.4-9.4-24.6 0-33.9l96.4-96.4-96.4-96.4c-9.4-9.4-9.4-24.6 0-33.9L54.3 103c9.4-9.4 24.6-9.4 33.9 0l136 136c9.5 9.4 9.5 24.6.1 34z"></path></svg>Join our Discord community](https://discord.com/invite/EEkFwbv)

## Have questions or feedback?

Follow and hit me up on Twitter [@q\_vault](https://twitter.com/q_vault) if you have any questions or comments. If I’ve made a mistake in the article be sure to [let me know](https://qvault.io/contact/) so I can get it corrected!