Implementing Red-Black Tree in Go — Part 1, The Search and Insert Operation

The binary search tree has one serious problem. They can transform into a linked list. The red-black tree solves this problem and ensures O(log(n)) time complexity for get, insert, and delete operations. To avoid this problem we can use self-balancing trees. The red-black tree is one of safe-balanced tree. It ensures that the depth between the right and left subtree does not differ by more than 1.

In addition to binary search properties, the red-black tree introduces additional properties:

1. The root of the red-black tree is black

2. The leaf is always black

3. The black node can have either a black or red node.

4. If the node is red, both children are black.

5. A simple path from a node to its descendant leaves always has the same number of black nodes.

   
   

We will also keep a separate instance per tree leaf. We could create an instance for each leaf, but that would be more memory-intensive. Also, we will create a comparator that will compare keys.

package util

import (
	"cmp"
)

type Comparator[T any] interface {
	Compare(a T, b T) int
}

type OrderedComparator[T cmp.Ordered] struct {
}

func (o OrderedComparator[T]) Compare(a T, b T) int {
	return cmp.Compare(a, b)
}

package redblacktree

import (
	"github.com/dmitriiantonov/go-ds/tree"
	"github.com/dmitriiantonov/go-ds/util"
)

type color int

const (
	black color = iota
	red
)

type TreeNode[K any, V any] struct {
	key    K
	value  V
	color  color
	parent *TreeNode[K, V]
	left   *TreeNode[K, V]
	right  *TreeNode[K, V]
}

func (node *TreeNode[K, V]) Key() K {
	return node.key
}

func (node *TreeNode[K, V]) Value() V {
	return node.value
}

type Tree[K any, V any] struct {
	leaf       *TreeNode[K, V]
	root       *TreeNode[K, V]
	comparator util.Comparator[K]
	len        int
}

func NewTree[K any, V any](comparator util.Comparator[K]) *Tree[K, V] {
	instance := new(Tree[K, V])
	instance.leaf = &TreeNode[K, V]{color: black}
	instance.root = instance.leaf
	instance.comparator = comparator
	return instance
}

func (tree *Tree[K, V]) Len() int {
	return tree.len
}

The difference from a binary search tree is that we end the search if the current node is leaf,

rather than nill.

The more interesting differences begin when inserting the node in our tree. We mark the new node as red and we can violate property 1 if the tree is empty, and property 4 because we can’t have two red nodes in the path. And, we must rebalance our tree.

First, we need to introduce the two operations used to balance the tree. These are rotating the tree to the left and to the right.

func (tree *Tree[K, V]) leftRotate(x *TreeNode[K, V]) {
	right := x.right
	x.right = right.left
	if right.left != tree.leaf {
		right.left.parent = x
	}
	right.parent = x.parent
	if x.parent == tree.leaf {
		tree.root = right
	} else if x == x.parent.left {
		x.parent.left = right
	} else {
		x.parent.right = right
	}
	right.left = x
	x.parent = right
}

func (tree *Tree[K, V]) rightRotate(x *TreeNode[K, V]) {
	left := x.left
	x.left = left.right
	if left.right != tree.leaf {
		left.right.parent = x
	}
	left.parent = x.parent
	if x.parent == tree.leaf {
		tree.root = left
	} else if x == x.parent.left {
		x.parent.left = left
	} else {
		x.parent.right = left
	}
	left.right = x
	x.parent = left
}

We have three cases in which we must fix violations in our tree. Also, we need to denote an uncle node, that is a sibling node’s parent.

1. If the uncle is a red node and the node’s parent is red we must make them black. Also, we must make their parent red and set it as the current node. If its parent is red, then we must fix the following violation.
2. If the node belongs to the left subtree and the node’s right child is red or if the node belongs to the right subtree and the node’s left child is red we must make the node’s parent current node and do the left rotation in the 1 case, or the right rotation in the second case. And then we come to case 3.
3. In the third case, we must make the current node’s parent black and its parent red, and then we must do a right rotation if the node’s parent is a left subtree, or a left rotation if the node’s parent is a right subtree.

I illustrated cases for the left subtree. For the right subtree, you must change the right rotation to the left rotation, and the left rotation to the right rotation.

func (tree *Tree[K, V]) Insert(key K, value V) {
    x := tree.lookup(key) // try to find a node by key

    if x != tree.leaf { // if we found it, update value and exit
       x.value = value
       return
    }
    // find the parent
    parent := tree.leaf
    current := tree.root

    for current != tree.leaf {
       parent = current

       switch tree.comparator.Compare(key, current.key) {
       case -1:
          current = current.left
       case 1:
          current = current.right
       }
    }
    // create a new node
    x = &TreeNode[K, V]{
       key:    key,
       value:  value,
       color:  red,
       parent: parent,
       left:   tree.leaf,
       right:  tree.leaf,
    }

    // send the node as child
    if parent == tree.leaf {
       tree.root = parent
    } else if tree.comparator.Compare(key, parent.key) < 0 {
       parent.left = x
    } else {
       parent.right = x
    }
    // do the rebalance
    tree.insertFixup(x)
}

func (tree *Tree[K, V]) insertFixup(x *TreeNode[K, V]) {
    for x.parent.color == red { // we must iterate only if the node's parent is red
       if x.parent == x.parent.parent.left { // if the node belongs to left subtree
          y := x.parent.parent.right // set uncle
          if y.color == red { // case 1
             x.parent.color = black
             y.color = black
             x.parent.parent.color = red // set 
             x = x.parent.parent
          } else {
             if x == x.parent.right { // case 2
                x = x.parent
                tree.leftRotate(x)
             }
             x.parent.color = black // case 3
             x.parent.parent.color = red
             tree.rightRotate(x.parent.parent)
          }
       } else { // if the node belongs to the right subtree
          y := x.parent.parent.left 
          if y.color == red { // case 1
             x.parent.color = black
             y.color = black
             x.parent.parent.color = red
             x = x.parent.parent
          } else {
             if x == x.parent.left { // case 2
                x = x.parent
                tree.rightRotate(x)
             }
             x.parent.color = black // case 3
             x.parent.parent.color = red
             tree.leftRotate(x.parent.parent)
          }
       }
    }
    tree.root.color = black // change the root color to black
}

In the second part of this article, we will consider the delete operation.