algorithms · leetcode · 2017-08-06 · yuex

The first time I get to know red black tree is quite a few years ago when I was reading CLRS' famous Introduction to Algorithms. But at that time their explaination of red black tree was very tedious and couter-intuitive. Perhaps that's because the original red black tree is not very neat by itself. But anyway, Sedgewick later discovered the left-leaning red black tree which leads to a very neat implementation. It's so neat that it actually can be added incrementally to your non-balanced tree implementation.

Let's try it out on Leetcode 315. This problem could be solved trivally by brute-force which results in \(O(n^2)\) time complexity. A better way is scan from right to left and insert integers into a BST while maintaining the number of nodes on the left of each node. If the BST is always balanced, the total time comlexity should be \(O(nlogn)\).

The non-balanced solution is given below

def countSmaller(self,ns):

    def new(val):
        node = TreeNode(val)
        node.size = 0
        node.count = 1
        return node

    def insert(root,node):
        if root == None:
            return node,node.size

        if node.val < root.val:
            root.size += 1
            root.left,size = insert(root.left,node)
        elif node.val > root.val:
            root.right,size = insert(root.right,node)
            size += (root.size + root.count)
        else: # ==
            root.count += 1
            size = root.size
        return root,size

    ret,root = [0 for _ in ns],None
    for i in xrange(len(ns)-1,-1,-1):
        root,size = insert(root,new(ns[i]))
        ret[i] = size
    return ret

Definitely, the worst-case time complexity is still \(O(n^2)\). Now let's modified above code incrementally to be a red black tree (left leaning) implementation. But first let's talk about some theories. This explaination is largely based on Sedgewick's video which is not same with the explaination from wikipedia where the color is taken as a property of a node instead an edge. The wikipedia's explaination is probably based on the original red black tree instead of the left leaning one. You can see at the end that the left leaning red black tree has much less cases and much neater code.

A red black tree is defined as binary tree. At first this definition may seem weird. But if you understand 2-3 tree where this type of red black tree is orginated from, following definition is quite natural.

  1. edges are either red or black
  2. a node can have at most 1 red edge to its children
  3. red edge of a node must be on the left
  4. every path from root to null has same number of black edges

From this definition, red edges occurs only at leaf nodes. Thus, we have only 3 different types of nodes. Two of them are leaves. There are also 3 operations related to red black tree

operation

To keep the BST balanced, we have to adjust the tree if insertion causes an unbalanced BST. Since we only have two different types of leaves, we will only have five different cases to insert a node. Here, we assume \(a \lt b \lt c\).

case

To get neat implementation, there are three key points. First we maintain the color of an edge at the destination node. Second we insert recursively. Third, the five cases could be reduce to four if we insert recursively as indicated in above picture. And only three of them need adjustment.

After some work, we can carry out the implementation.

def counSmaller(self,ns):

    def new(val):
        node = TreeNode(val)
        node.size = 0
        node.count = 1
        node.color = True
        return node

    def rotate_left(root):
        right = root.right
        root.color,right.color = right.color,root.color
        root.right,right.left = right.left,root
        right.size += (root.size + root.count)
        return right

    def rotate_right(root):
        left = root.left
        root.color,left.color = left.color,root.color
        root.left,left.right = left.right,root
        root.size -= (left.size + left.count)
        return left

    def flip(root):
        root.left.color = False
        root.right.color = False
        root.color = True
        return root

    def red(node):
        return node != None and node.color

    def insert(root,node):
        if root == None:
            return node,node.size

        if node.val == root.val:
            root.count += 1
            return root,root.size

        if node.val < root.val:
            root.size += 1
            root.left,size = insert(root.left,node)
        else: # node.val > root.val:
            root.right,size = insert(root.right,node)
            size += (root.size + root.count)

        if red(root.right) and not red(root.left):
            root = rotate_left(root)
        if red(root.left) and red(root.left.left):
            root = rotate_right(root)
        if red(root.left) and red(root.right):
            root = flip(root)

        return root,size


    ret,root = [0 for _ in ns],None
    for i in xrange(len(ns)-1,-1,-1):
        root,size = insert(root,new(ns[i]))
        ret[i] = size
    return ret

As you can see, except those helper functions, only 6 lines of code are added at the end of insert(). Although I shortcut the equal case, it is nothing substancial. The helper functions are also very easy to implement. So probably you can try to give a red black tree implement during a 45-mintue interview. It's quite feasible. But due to the test cases design of this problem, the red black tree implementation is acutally one time slower in terms of total execution time than the non-balance implementaion.