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 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$$. 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.