Red-Black Trees Part 4 - Optimization

By John Lenz. December 27, 2015.

This page is part four of the article series You could have invented red-black trees!.

So far, the pseudocode and discussion of red-black trees focused on correctness. How is this pseudocode translated into an actual implementation? For example, red-black trees are used in the Linux kernel and you can see their implementation in Linux 4.3 (the latest release as I write this) in rbtree.h and rbtree.c. with some documentation in rbtree.txt.

Tail-Recursion

Thinking about the recursive procedures for insert and remove, one place where they will do useless work is when nothing needs to be fixed. For example, during remove if we happen to remove a red node we saw that nothing needs to be fixed. In this situation, the recursive procedure will descend down the tree, remove the node, and then ascend back up the tree doing nothing (besides executing some if statements that will always be false).

We therefore might ask ourselves if we can pre-emptively perform the fixes, rotations, and color changes on the way down the tree so that nothing ever needs to be done on the way back up the tree. For example, if on the way down the tree we could perform some flips or color changes so as to always guarantee that we remove a red node, we could avoid ascending back up the tree. If we could do so, we can then make the procedure tail-recursive since nothing happens after recursion returns, and then we can either manually convert it to an efficient loop or rely on the compiler to perform this optimization.

The answer happens to be yes. Consider remove where we want to guarantee on our way down the tree that we will eventually remove a red node. The root is black so at the beginning of insert you can change the root to red. (This can possibly create a red-edge between the root and one of its children.) As we descend down the tree we will drag along the red color (and also dragging along the red-edge if it exists) by performing flips and color changes. This is sometimes called top-down removal.

If you implement this, you might be surprised that it sometimes performs worse. What happens is that sometimes this top-down remove procedure will perform many more rotations than the bottom-up remove discussed on the previous page which fixes the problems on the way back up the tree. While each rotation is \(O(1)\), in practice more rotations will lead to more memory operations and a slower procedure. Briefly, the bottom-up remove described on the previous page guarantees that if it performs a rotation, then the double-black node has been removed and no more fixing occurs. Indeed, the fixing during the bottom-up remove is a sequence of pull_black operations followed by a single fix_* procedure which rotates and removes the double-black color. In contrast, the top-down remove can perform many rotations as it is pushing the red color down the tree.

A similar story happens for insert: it is possible to implement top-down insert where you perform rotations on the way down the tree so that nothing needs to be fixed on the way back up the tree (and therefore the ascent can be skipped by tail-recursion). But, the upper bound on the number of possible rotations is larger than for the bottom-up procedure.

Parent Pointers

We are back to square one. We still want to descend down the tree, perform the insert or removal, and then fix the problem (red-edge or double-black node) on the way back up the tree. But we still sometimes have useless work on the way back up the tree (for example when we remove a red node). Sometimes we actually must fix all the way up the tree so we can't avoid the ascent. What we want is a way of ascending up the tree and aborting once we have fixed the problem.

To do so, we can add a parent pointer to each node in the tree. If you like recursion, we could then phrase the fixing ascent as a tail-recursive procedure which starts at the leaf and recurses to the parent. The base case is the root. But at this point it makes more sense to use loops. We will have two loops: one for descending down the tree and one for ascending up the tree. Either loop can abort early by just returning.

The first task is to update flip and rotate functions to keep parent pointer's and the root updated. For example, when rotating at the root the new root must be set. When rotating at a node that is not the root, the parent of the node must have its child pointer be updated to point to the new result of the rotate. The flip and rotate functions then no longer need to return anything. I'm not going to give pseudocode for these updated rotate and flip utilities, they are straightforward modifications.

Iterative Insert

The first task of insert is to descend down the tree according to the binary search tree property. Converted to a loop, this looks like the following pseudocode:

# Descend down the tree
g = root
while true:
    if newKey < g.key:
        if g.left = nil:
            g.left = new red node with newKey,newVal
            g = g.left
            break
        else:
            g = g.left

    elif newKey > g.key:
        if g.right = nil:
            g.right = new red node with newKey, newVal
            g = g.right
            break
        else:
            g = g.right
    else:
        g.val = newVal
        return

This loop ends with g the newly added node, so the next task is to ascend back up the tree. The recursive insert procedure ascends back up the tree calling fix_red_edge on each node. We want the body of the fix_red_edge function to become the body of a loop which ascends up the tree, which is the following procedure (you should look back at the fix_red_edge procedure on the insert page). Note that the rotate and flip utilities no longer return anything since they can fix the parent directly.

while g != nil:
    if g.left.color = red and g.right.color = red:
        push_black(g)
        g = g.parent
        continue loop

    if g.left.color = red and g.left.right.color = red:
        rotate_left(g.left)
    if g.left.color = red and g.left.left.color = red:
        flip_right(g)

    if g.right.color = red and g.right.left.color = red:
         rotate_right(g.right)
    if g.right.color = red and g.right.right.color = red:
        flip_left(g)

    g = g.parent

The above code ascends up the tree, performing the exact same rotates and flips as the recursive procedure (thus we know it is correct). At the moment, the loop runs all the way up to the root, but the point of adding parent pointers is that we could abort the ascent early. There are several optimizations that we can perform here:

• Looking back at insert, we can see that if we perform a rotate or flip, the red-edge is removed completely (go back through the figures). Thus after flipping, we can break out of the loop since the tree is a red-black tree. Note that if we rotate because g.left and g.left.right are red, we will immediately flip and so can also break.

• It is possible to skip over some nodes during the ascent. Consider that we are currently located at a red node. (For example, this happens at the very beginning of the ascent since we add a new node in red.) If the parent of g is black, we can immediately abort the ascent since there is no red-edge. If the parent of g is red, we have a red-edge and can jump directly to the grandparent of g since that is the node from which we will fix the red-edge. Finally, if the parent of g is the root, we can just abort since we always set the root to black.

• It is now impossible to be located at a black node during the ascent. We start at a red node, and the only way the loop continues is after a push_black which sets g's color to red. The loop then continues, and if g's parent is red we have a red-edge and jump directly to g's grandparent. This means that every iteration of the loop will perform some operation, either pushing black and continuing or rotating/flipping and breaking out of the loop. For debugging, we can therefore add assertion that the loop body either continues due to push_black with a red parent, or breaks.

We therefore have the following loop:

while true:
    # g.color is red

    if g = root or g.parent = root or g.parent.color = black:
        break
    g = g.parent.parent

    if g.left.color = red and g.right.color = red:
        push_black(g)
        continue loop

    if g.left.color = red and g.left.right.color = red:
        rotate_left(g.left)
    if g.left.color = red and g.left.left.color = red:
        flip_right(g)
        break

    if g.right.color = red and g.right.left.color = red:
        rotate_right(g.right)
    if g.right.color = red and g.right.right.color = red:
        flip_left(g)
        break

    throw "Fail: Should never occur!"

root.color = black

There is one final optimization that we can perform here: many of these if statements are redundant. If we have not aborted the loop, we actually know the location of the red-edge based on the direction of ascent. So instead of just setting g = g.parent.parent, we will track the direction of ascent and then use that to remove many of the if statements from the code, leaving only a few color checks. I won't show pseudocode for this, because after this optimization we have essentially arrived at the Linux kernel implementation.

The ascent loop starts on line 102. The variables are a little different, but you can clearly see the operations. First, on lines 110 to 114 it is checking if the parent is black and if so aborting. push_black and continue start on line 120 and is called "Case 1" (the diagrams in the comments use upper case letters for black nodes and lower case letters for red). The rotate left when g.left.right is red is called "Case 2" and is located on line 144, and then the flip right is called "Case 3" and is on line 168. After "Case 3", the code breaks out of the loop. Note that there is no if statement around Case 3! This is because we are in the if statement on line 119 which says that we ascended from the left and know that we are not pushing black so must be flipping right. The reflected operations for when you ascend from the right starts on line 185.

Iterative Remove

If you scroll down in rbtree.c, the ascent loop which fixes the tree after removing a node starts on line 233. Similar to insert, we can arrive at this loop via a sequence of optimizations performed on the recursive remove procedure we discussed before. If you look back at the recursive remove pseudocode, we called the fix_doubleblack function on each node during the ascent up the tree. Thus, fix_doubleblack will become the body of the ascent loop. Like before, there are several optimizations we can perform. I will only outline them, you should try it yourself to compare the fix_doubleblack function and the rbtree.c loop.

First, we don't need to store the double-black color in the tree itself. Since we are going to abort the loop once the double-black color no longer exists, we can just always guarantee that the node we are currently located at is the double-black node. This will require some updates to the pull, push, rotate, and flip utility functions, but this is OK since we are going to manually inline those utility functions directly where they are used. If you are looking at rbtree.c, this is described in the comment starting on line 235. The comment states that the node is black and that paths though node have one fewer black node. Translated into our language, this means node is the double-black node even though in memory the color of node is black.

Many of the if statements in fix_doubleblack and friends are redundant depending on the direction of ascent. Despite this, you can still see rbtree.c still performs the same operations in the same order as fix_doubleblack. For example, the if statement on line 242 is true when the double-black node is to the left of the variable parent. In the diagrams, the double-black node is N. "Case 1" is when the sibling of N is red, which inside fix_left_doubleblack caused us to flip left before calling fix_left_doubleblack_right_black. "Case 1" carries out the fix, and then falls down to the code starting on line 262. Therefore, the fix_left_doubleblack_right_black function starts on line 262. For comparison, I copied here the fix_left_doubleblack_right_black function that we wrote:

def fix_left_doubleblack_right_black(z):
    assert(z.left.color == double-black and z.right.color == black)
    pull_black(z)
    if z.right.left.color == red and z.right.right.color == black:
        z.right = rotate_right(z.right)
    if z.right.right.color == red:
        z = flip_left(z)
        push_black(z)
    return z

"Case 2" on line 266 is when a pull_black occurs with both right grandchildren black (the grandchildren are Sl and Sr in the figure). After pulling, we might have moved the double-black color up the tree and if so we can continue the loop. If the double-black color was removed (which happens when p is red), we can just abort the loop.

"Case 3" on line 293 is when z.right.left is red and z.right.right is black, where fix_left_doubleblack_right_black rotated right. "Case 4" on line 317 is when z.right and z.right.right are black, where we flip left and the push black. The code for "Case 4" inlines and combines the code for the flip left and push black. After performing these rotations and flips, the double-black node has been removed so we can just break on line 337.

The code starting on line 339 then fixes the situation when the double-black node is to the right.