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.
- edges are either red or black
- a node can have at most 1 red edge to its children
- red edge of a node must be on the left
- 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.