Red-black.py

class RedBlackTree:
    #Red/Black tree implementation based on 
    #Algorithms in C++, Sedgewick
    #Introduction To Algorithms Cormen, Thomas H. / Leiserson, Charles E . / Rivest, Ronald L . The MIT Press 07/1990
    # NOTE : LOOK AT END OF FILE TO SEE DIFFERENCES IN TRAVERSAL IDIOMS
    
    red,black = range(2)

    # double underscores induce name mangling to make methods private
    def __copy(self,node):
        self.__left  = node.__left
        self.__right = node.__right
        self.__val   = node.__val
        self.__color = node.__color
        
    def __RBinsertLeft(self,val,sw):
        if (self.__left == None):
            self.__left = RedBlackTree(val)
        else:
            self.__left.__RBinsert(val,sw)
            
    def __RBinsertRight(self,val,sw):
        if (self.__right == None):
            self.__right = RedBlackTree(val)
        else:
            self.__right.__RBinsert(val,sw)

    def __rotLeft(self):
       if (self.__right == None):
           return None
           
       # make the changes in a copy of current node __self
       # on return, the caller will copy in 'me' to current node
       me = RedBlackTree()
       me.__copy(self)
       x           = me.__right
       me.__right  = x.__left
       x.__left    = me
       return   x

    def __rotRight(self):
       if (self.__left == None):
           return None
           
       # make the changes in a copy of current node __self
       # on return, the caller will copy in 'me' to current node
       me = RedBlackTree()
       me.__copy(self)
       x           = me.__left;
       me.__left   = x.__right;
       x.__right   = me;
       return x

    def __RBinsert(self,val,sw):
        # if current node is a 4 node, split it by flipping its colors
        # if both children of this node are red, change this one to red
        # and the children to black
        left  = self.__left
        right = self.__right
        if ((left != None)and(left.__color==RedBlackTree.red)and(right != None)and(right.__color==RedBlackTree.red)):
            self.__color  = RedBlackTree.red
            left.__color  = RedBlackTree.black
            right.__color = RedBlackTree.black
        
        # go left or right depending on key relationship
        if (val < self.__val):
            # recursively insert
            self.__RBinsertLeft(val,0)
            
            # on the way back up check if a rotation is needed
            # if search path has two red links with same orientation
            # do a single rotation and flip the color bits
            left = self.__left
            if ((self.__color == RedBlackTree.red)and(left != None)and(left.__color == RedBlackTree.red)and(sw == 1)):
                x = self.__rotRight()
                if (x != None):
                    self.__copy(x)
           
           # flip the color bits
            left = self.__left
            if (left != None):
                ll = left.__left
                if (ll != None):
                    if ((left.__color == RedBlackTree.red)and(ll.__color == RedBlackTree.red)):
                        x = self.__rotRight()
                        if (x != None):
                            self.__copy(x)
                        self.__color = RedBlackTree.black
                        right = self.__right
                        if (right != None):
                           right.__color = RedBlackTree.red
                           
        else:
            self.__RBinsertRight(val,1)
            
            # on the way back up check if a rotation is needed
            # if search path has two red links with same orientation
            # do a single rotation and flip the color bits
            right = self.__right
            if ((self.__color == RedBlackTree.red)and(right != None)and(right.__color == RedBlackTree.red)and(sw == 0)):
                x = self.__rotLeft()
                if (x != None):
                    self.__copy(x)
           
           # flip the color bits
            right = self.__right
            if (right != None):
                rr = right.__right
                if (rr != None):
                    if ((right.__color == RedBlackTree.red)and(rr.__color == RedBlackTree.red)):
                        x = self.__rotLeft()
                        if (x != None):
                            self.__copy(x)
                        self.__color = RedBlackTree.black
                        left = self.__left
                        if (left != None):
                           left.__color = RedBlackTree.red


# ============================================================
# public methods
# ============================================================
    # constructor
    def __init__(self,val = None):
        self.__left      = None
        self.__right     = None
        self.__val       = val
        self.__color     = RedBlackTree.red

    # to string
    def __str__(self):
        return "[" + str(self.__val) + "," + str(self.__color) + "]"
    
    # accessor
    def val(self):
        return self.__val
        
    # accessor
    def color(self):
        return self.__color        

    # recursive 'find'
    def find(self,key):
        result = None
        if (key == self.__val):
            result = self
        elif (key < self.__val):
            if (self.__left != None):
                result = self.__left.find(key)
        else:
            if (self.__right != None):
                result = self.__right.find(key)
        return result

    # inorder traversal using a class as the visitor
    def inorderClass(self,visitor,depth):
        if self.__left != None:
            self.__left.inorderClass(visitor,depth+1)
        visitor.visit(self,depth)
        if self.__right != None:
            self.__right.inorderClass(visitor,depth+1)
            

    # inorder traversal using a function as the visitor
    def inorderFunction(self,visit,depth):
        if self.__left != None:
            self.__left.inorderFunction(visit,depth+1)
        visit(self,depth)
        if self.__right != None:
            self.__right.inorderFunction(visit,depth+1)
            
    # main public 'insert' function        
    def insert(self,val):
        self.__RBinsert(val,0);
        self.__color = RedBlackTree.black

# define a visitor class if you need the visitor to have state     
class NodeVisitor:
    def __init__(self):
        pass
        
    def visit(self,node,depth):
        if (node.val() != None):
            print ( "(" + str(node.val()) + ":"+ str(node.color()) + ":" + str(depth) + "),"),
            
# just use a function if your visitor doesn't need state
def visit(node,depth):
    if (node.val() != None):
        print ("(" + str(node.val()) + ":"+ str(node.color()) + ":" + str(depth) + "),"),
    
# ==================================
# test program
# ==================================
if __name__ == "__main__":
    nodelist = [11,4,8,14,17,6,9,7,16,15,13,5,19,18,12,10,3,20,1]
    root  = RedBlackTree(2)
    for n in nodelist:
        root.insert(n)
        
    # pass a class instance
    v = NodeVisitor()
    print ("Python C ="),    
    root.inorderClass(v,0)
    print

    # use a named function
    print ("Python F ="),    
    root.inorderFunction(visit,0)
    print
    
    # do a find
    t = root.find(16)
    print (str(t))