Red-Black Tree Implementation in Python

Question

I'm wondering if there are better practices, and optimizations I could make on my code, to do things easier/faster.
I have added comments, in many places in the code, to make understanding easier.

import sys
class Node():
 def __init__(self, data):
 self.data = data # holds the key
 self.parent = None #pointer to the parent
 self.left = None # pointer to left child
 self.right = None #pointer to right child
 self.color = 1 # 1 . Red, 0 . Black
# class RedBlackTree implements the operations in Red Black Tree
class RedBlackTree():
 def __init__(self, List=None):
 self.TNULL = Node(0)
 self.TNULL.color = 0
 self.TNULL.left = None
 self.TNULL.right = None
 self.root = self.TNULL
 if List:
 for val in List:
 self.insert(val)
 def __pre_order_helper(self, node):
 if node != TNULL:
 sys.stdout.write(node.data + " ")
 self.__pre_order_helper(node.left)
 self.__pre_order_helper(node.right)
 def __in_order_helper(self, node):
 if node != TNULL:
 self.__in_order_helper(node.left)
 sys.stdout.write(node.data + " ")
 self.__in_order_helper(node.right)
 def __post_order_helper(self, node):
 if node != TNULL:
 self.__post_order_helper(node.left)
 self.__post_order_helper(node.right)
 sys.stdout.write(node.data + " ")
 def __search_tree_helper(self, node, key):
 if node == TNULL or key == node.data:
 return node
 if key < node.data:
 return self.__search_tree_helper(node.left, key)
 return self.__search_tree_helper(node.right, key)
 # fix the rb tree modified by the delete operation
 def __fix_delete(self, x):
 while x != self.root and x.color == 0:
 if x == x.parent.left:
 s = x.parent.right
 if s.color == 1:
 # case 3.1
 s.color = 0
 x.parent.color = 1
 self.left_rotate(x.parent)
 s = x.parent.right
 if s.left.color == 0 and s.right.color == 0:
 # case 3.2
 s.color = 1
 x = x.parent
 else:
 if s.right.color == 0:
 # case 3.3
 s.left.color = 0
 s.color = 1
 self.right_rotate(s)
 s = x.parent.right
 # case 3.4
 s.color = x.parent.color
 x.parent.color = 0
 s.right.color = 0
 self.left_rotate(x.parent)
 x = self.root
 else:
 s = x.parent.left
 if s.color == 1:
 # case 3.1
 s.color = 0
 x.parent.color = 1
 self.right_rotate(x.parent)
 s = x.parent.left
 if s.left.color == 0 and s.right.color == 0:
 # case 3.2
 s.color = 1
 x = x.parent
 else:
 if s.left.color == 0:
 # case 3.3
 s.right.color = 0
 s.color = 1
 self.left_rotate(s)
 s = x.parent.left 
 # case 3.4
 s.color = x.parent.color
 x.parent.color = 0
 s.left.color = 0
 self.right_rotate(x.parent)
 x = self.root
 x.color = 0
 def __rb_transplant(self, u, v):
 if u.parent == None:
 self.root = v
 elif u == u.parent.left:
 u.parent.left = v
 else:
 u.parent.right = v
 v.parent = u.parent
 def __delete_node_helper(self, node, key):
 # find the node containing key
 z = self.TNULL
 while node != self.TNULL:
 if node.data == key:
 z = node
 if node.data <= key:
 node = node.right
 else:
 node = node.left
 if z == self.TNULL:
 print("Couldn't find key in the tree")
 return
 y = z
 y_original_color = y.color
 if z.left == self.TNULL:
 x = z.right
 self.__rb_transplant(z, z.right)
 elif (z.right == self.TNULL):
 x = z.left
 self.__rb_transplant(z, z.left)
 else:
 y = self.minimum(z.right)
 y_original_color = y.color
 x = y.right
 if y.parent == z:
 x.parent = y
 else:
 self.__rb_transplant(y, y.right)
 y.right = z.right
 y.right.parent = y
 self.__rb_transplant(z, y)
 y.left = z.left
 y.left.parent = y
 y.color = z.color
 if y_original_color == 0:
 self.__fix_delete(x)
 
 # fix the red-black tree
 def __fix_insert(self, k):
 while k.parent.color == 1: # while k's parent is red
 if k.parent == k.parent.parent.right: # if k's parent is the right child of k's grandparent, then the uncle is the left child of k's grandparent
 u = k.parent.parent.left # uncle
 if u.color == 1: # if k's parent and uncle are both red ---> recolor both uncle and the parent to black, and the grandparent to red.
 # case 3.1
 u.color = 0
 k.parent.color = 0
 k.parent.parent.color = 1
 k = k.parent.parent
 else: # if k's parent is red, and uncle is black
 if k == k.parent.left:
 # case 3.2.2, if parent is the right child of grandparent and k is the left child of parent (Right-Left) ---> Perform Right Rotation on Parent, and it becomes case 3.2.1
 k = k.parent
 self.right_rotate(k)
 # case 3.2.1, if parent is the right child of grandparent and k is the right child of parent {Right-Right} ---> Perform Left Rotation on grandparent, and recolor k's new parent to black, and it's sibling to red
 k.parent.color = 0
 k.parent.parent.color = 1
 self.left_rotate(k.parent.parent)
 else: # if k's parent is the left child of k's grandparent, then the uncle is the right child of k's grandparent
 u = k.parent.parent.right # uncle
 if u.color == 1:
 # mirror case 3.1
 u.color = 0
 k.parent.color = 0
 k.parent.parent.color = 1
 k = k.parent.parent 
 else:
 if k == k.parent.right:
 # mirror case 3.2.2
 k = k.parent
 self.left_rotate(k)
 # mirror case 3.2.1
 k.parent.color = 0
 k.parent.parent.color = 1
 self.right_rotate(k.parent.parent)
 if k == self.root:
 break
 self.root.color = 0
 def __print_helper(self, node, indent, last):
 # print the tree structure on the screen
 if node != self.TNULL:
 sys.stdout.write(indent)
 if last:
 sys.stdout.write("R----")
 indent += " "
 else:
 sys.stdout.write("L----")
 indent += "| "
 s_color = "RED" if node.color == 1 else "BLACK"
 print(str(node.data) + "(" + s_color + ")")
 self.__print_helper(node.left, indent, False)
 self.__print_helper(node.right, indent, True)
 
 # Pre-Order traversal
 # Node.Left Subtree.Right Subtree
 def preorder(self):
 self.__pre_order_helper(self.root)
 # In-Order traversal
 # left Subtree . Node . Right Subtree
 def inorder(self):
 self.__in_order_helper(self.root)
 # Post-Order traversal
 # Left Subtree . Right Subtree . Node
 def postorder(self):
 self.__post_order_helper(self.root)
 # search the tree for the key k
 # and return the corresponding node
 def searchTree(self, k):
 return self.__search_tree_helper(self.root, k)
 # find the node with the minimum key
 def minimum(self, node):
 while node.left != self.TNULL:
 node = node.left
 return node
 # find the node with the maximum key
 def maximum(self, node):
 while node.right != self.TNULL:
 node = node.right
 return node
 # find the successor of a given node
 def successor(self, x):
 # if the right subtree is not None,
 # the successor is the leftmost node in the
 # right subtree
 if x.right != self.TNULL:
 return self.minimum(x.right)
 # else it is the lowest ancestor of x whose
 # left child is also an ancestor of x.
 y = x.parent
 while y != self.TNULL and x == y.right:
 x = y
 y = y.parent
 return y
 # find the predecessor of a given node
 def predecessor(self, x):
 # if the left subtree is not None,
 # the predecessor is the rightmost node in the 
 # left subtree
 if (x.left != self.TNULL):
 return self.maximum(x.left)
 y = x.parent
 while y != self.TNULL and x == y.left:
 x = y
 y = y.parent
 return y
 # rotate left at node x
 def left_rotate(self, x):
 y = x.right
 x.right = y.left
 if y.left != self.TNULL:
 y.left.parent = x
 y.parent = x.parent
 if x.parent == None:
 self.root = y
 elif x == x.parent.left:
 x.parent.left = y
 else:
 x.parent.right = y
 y.left = x
 x.parent = y
 # rotate right at node x
 def right_rotate(self, x):
 y = x.left
 x.left = y.right
 if y.right != self.TNULL:
 y.right.parent = x
 y.parent = x.parent
 if x.parent == None:
 self.root = y
 elif x == x.parent.right:
 x.parent.right = y
 else:
 x.parent.left = y
 y.right = x
 x.parent = y
 # insert the key to the tree in its appropriate position
 # and fix the tree
 def insert(self, key):
 # Ordinary Binary Search Insertion
 node = Node(key)
 node.parent = None
 node.data = key
 node.left = self.TNULL
 node.right = self.TNULL
 node.color = 1 # new node must be red
 y = None
 x = self.root
 while x != self.TNULL:
 y = x
 if node.data < x.data:
 x = x.left
 else:
 x = x.right
 # y is parent of x
 node.parent = y
 if y == None:
 self.root = node
 elif node.data < y.data:
 y.left = node
 else:
 y.right = node
 # if new node is a root node, simply return
 if node.parent == None:
 node.color = 0
 return
 # if the grandparent is None, simply return
 if node.parent.parent == None:
 return
 # Fix the tree
 self.__fix_insert(node)
 def get_root(self):
 return self.root
 # delete the node from the tree
 def delete_node(self, data):
 self.__delete_node_helper(self.root, data)
 # print the tree structure on the screen
 def pretty_print(self):
 self.__print_helper(self.root, "", True)
if __name__ == "__main__":
 bst = RedBlackTree([1,2,3,4,5,6])

Answer 1

Docstrings

You are using comments instead of docstrings, which is not a good tone, as I think. You can read more about docstring in PEP 257. Also it's not necessary to write class name in docstring. For example, you can write Implements the operations in Red Black Tree instead of class RedBlackTree implements the operations in Red Black Tree.

Type hinting

Your code doesn't have type hinting at all. It's not really a bad thing, but it's always better to do type hinting. You can read more about why you may want to do it in PEP 484 and how to do it here.

Double underscores in method names

It's not really necessary to write __ before the method name, because _ is already enough. But it's a matter of preference.

Node colors

I'd use enum instead of 0/1 for colors. Your code will be more understandable if you write self.color = Color.Red instead of self.color = 1.

Node.data

Data is very wide and uninformative name. If data holds key, why don't you name it key?

sys.stdout

I suppose that there's no need to overcomplicate your code by using sys.stdout.write instead of print, which also prints to stdout.

RedBlackTree.pretty_print()

You should definitely read about dunders. They will allow you to integrate your class with default python methods. In this particular case dunders will allow you to use print(some_tree) instead of some_tree.pretty_print().

Long methods

Some of your methods are really long and difficult to read. For example, __fix_delete. Also comments like # case 3.1 are uninformative. I guess you should extract each case in separate function and write a little bit more about each case in docstrings.

Style comments

I think that every Python programmer should use pylint. It's a good tool to find small issues like too long lines, absence of spaces after commas, etc. It's really helpful, but try not to become pylint maniac - it's not always necessary to achieve 10/10.

• Reinderien: Re. _ is a matter of preference - it's more than that. Single- and double-underscored methods follow well-defined conventions - single underscores are for "weakly-enforced private variables", and double underscores are for name mangling. Read stackoverflow.com/questions/1301346