Implementing a Red-Black Tree in Golang

Red-Black Trees are a type of self-balancing binary search tree. Each node in the tree has an extra attribute for denoting the color of the node, either red or black. The tree satisfies the following properties:

1. Every node is either red or black.

2. The root is always black.

3. If a node is red, then both its children are black.

4. Every path from a node to its descendant NIL nodes has the same number of black nodes.

5. New insertions are always red.

In this blog post, we'll walk through the process of implementing a Red-Black Tree in Golang.

Node Structure

First, let's define the structure of a node:

type Color int

const (
    RED   Color = iota
    BLACK
)

type Node struct {
    Key    int
    Color  Color
    Parent *Node
    Left   *Node
    Right  *Node
}

Rotations

Rotations are fundamental operations that help in maintaining the Red-Black Tree properties. There are two types of rotations: left rotation and right rotation.

Left Rotation

func (n *Node) rotateLeft() *Node {
    rChild := n.Right
    n.Right = rChild.Left
    if rChild.Left != nil {
        rChild.Left.Parent = n
    }
    rChild.Parent = n.Parent
    if n.Parent == nil {
        // n is root
        root = rChild
    } else if n == n.Parent.Left {
        n.Parent.Left = rChild
    } else {
        n.Parent.Right = rChild
    }
    rChild.Left = n
    n.Parent = rChild
    return rChild
}

Right Rotation

func (n *Node) rotateRight() *Node {
    lChild := n.Left
    n.Left = lChild.Right
    if lChild.Right != nil {
        lChild.Right.Parent = n
    }
    lChild.Parent = n.Parent
    if n.Parent == nil {
        // n is root
        root = lChild
    } else if n == n.Parent.Right {
        n.Parent.Right = lChild
    } else {
        n.Parent.Left = lChild
    }
    lChild.Right = n
    n.Parent = lChild
    return lChild
}

Insertion

Insertion in a Red-Black Tree involves two steps:

1. Regular BST insertion.

2. Fixing the Red-Black Tree properties.

func (tree *RBTree) Insert(key int) {
    newNode := &Node{Key: key, Color: RED}
    if tree.Root == nil {
        tree.Root = newNode
    } else {
        currentNode := tree.Root
        var parentNode *Node
        for currentNode != nil {
            parentNode = currentNode
            if newNode.Key < currentNode.Key {
                currentNode = currentNode.Left
            } else {
                currentNode = currentNode.Right
            }
        }
        newNode.Parent = parentNode
        if newNode.Key < parentNode.Key {
            parentNode.Left = newNode
        } else {
            parentNode.Right = newNode
        }
    }
    tree.fixInsert(newNode)
}

Fixing the Tree

After insertion, the Red-Black Tree properties might be violated. We need to fix these violations.

func (tree *RBTree) fixInsert(n *Node) {
    parentNode := n.Parent

    // While the parent node is red, we need to restructure and/or recolor.
    for parentNode != nil && parentNode.Color == RED {
        grandparentNode := parentNode.Parent

        // Parent is the left child of the grandparent.
        if parentNode == grandparentNode.Left {
            uncle := grandparentNode.Right

            // Case 1: Uncle is red. Just recolor.
            if uncle != nil && uncle.Color == RED {
                grandparentNode.Color = RED
                parentNode.Color = BLACK
                uncle.Color = BLACK
                n = grandparentNode
            } else {
                // Case 2: Node is the right child. Left rotation on parent.
                if n == parentNode.Right {
                    n = parentNode
                    tree.rotateLeft(n)
                }

                // Case 3: Node is the left child. Right rotation on grandparent.
                parentNode.Color = BLACK
                grandparentNode.Color = RED
                tree.rotateRight(grandparentNode)
            }
        } else { // Parent is the right child of the grandparent.
            uncle := grandparentNode.Left

            // Case 1: Uncle is red. Just recolor.
            if uncle != nil && uncle.Color == RED {
                grandparentNode.Color = RED
                parentNode.Color = BLACK
                uncle.Color = BLACK
                n = grandparentNode
            } else {
                // Case 2: Node is the left child. Right rotation on parent.
                if n == parentNode.Left {
                    n = parentNode
                    tree.rotateRight(n)
                }

                // Case 3: Node is the right child. Left rotation on grandparent.
                parentNode.Color = BLACK
                grandparentNode.Color = RED
                tree.rotateLeft(grandparentNode)
            }
        }
        parentNode = n.Parent
    }

    // Ensure the root is always black.
    tree.Root.Color = BLACK
}

The fixInsert function will involve checking the color of the uncle node and deciding whether to perform a rotation or simply recolor the nodes.

Conclusion

This is a basic introduction to implementing a Red-Black Tree in Golang. There are many more details to consider, especially when it comes to deletion and the associated fix-up procedures. However, with the foundation provided here, you can further explore and complete the Red-Black Tree implementation.