1:#/usr/bin/env python
     2:class RedBlackTree:
     3:    #Red/Black tree implementation based on 
     4:    #Algorithms in C++, Sedgewick
Py   5:    #Introduction To Algorithms Cormen, Thomas H. / Leiserson, Charles E . / Rivest, Ronald L . The MIT Press 07/1990
     6:    # NOTE : LOOK AT END OF FILE TO SEE DIFFERENCES IN TRAVERSAL IDIOMS
     7:    
     8:    red,black = range(2)
     9:
Py  10:    # double underscores induce name mangling to make methods private
    11:    def __copy(self,node):
    12:        self.__left  = node.__left
    13:        self.__right = node.__right
    14:        self.__val   = node.__val
Py  15:        self.__color = node.__color
    16:        
    17:    def __RBinsertLeft(self,val,sw):
    18:        if (self.__left == None):
    19:            self.__left = RedBlackTree(val)
Py  20:        else:
    21:            self.__left.__RBinsert(val,sw)
    22:            
    23:    def __RBinsertRight(self,val,sw):
    24:        if (self.__right == None):
Py  25:            self.__right = RedBlackTree(val)
    26:        else:
    27:            self.__right.__RBinsert(val,sw)
    28:
    29:    def __rotLeft(self):
Py  30:       if (self.__right == None):
    31:           return None
    32:           
    33:       # make the changes in a copy of current node __self
    34:       # on return, the caller will copy in 'me' to current node
Py  35:       me = RedBlackTree()
    36:       me.__copy(self)
    37:       x           = me.__right
    38:       me.__right  = x.__left
    39:       x.__left    = me
Py  40:       return   x
    41:
    42:    def __rotRight(self):
    43:       if (self.__left == None):
    44:           return None
Py  45:           
    46:       # make the changes in a copy of current node __self
    47:       # on return, the caller will copy in 'me' to current node
    48:       me = RedBlackTree()
    49:       me.__copy(self)
Py  50:       x           = me.__left;
    51:       me.__left   = x.__right;
    52:       x.__right   = me;
    53:       return x
    54:
Py  55:    def __RBinsert(self,val,sw):
    56:        # if current node is a 4 node, split it by flipping its colors
    57:        # if both children of this node are red, change this one to red
    58:        # and the children to black
    59:        left  = self.__left
Py  60:        right = self.__right
    61:        if ((left != None)and(left.__color==RedBlackTree.red)and(right != None)and(right.__color==RedBlackTree.red)):
    62:            self.__color  = RedBlackTree.red
    63:            left.__color  = RedBlackTree.black
    64:            right.__color = RedBlackTree.black
Py  65:        
    66:        # go left or right depending on key relationship
    67:        if (val < self.__val):
    68:            # recursively insert
    69:            self.__RBinsertLeft(val,0)
Py  70:            
    71:            # on the way back up check if a rotation is needed
    72:            # if search path has two red links with same orientation
    73:            # do a single rotation and flip the color bits
    74:            left = self.__left
Py  75:            if ((self.__color == RedBlackTree.red)and(left != None)and(left.__color == RedBlackTree.red)and(sw == 1)):
    76:                x = self.__rotRight()
    77:                if (x != None):
    78:                    self.__copy(x)
    79:           
Py  80:           # flip the color bits
    81:            left = self.__left
    82:            if (left != None):
    83:                ll = left.__left
    84:                if (ll != None):
Py  85:                    if ((left.__color == RedBlackTree.red)and(ll.__color == RedBlackTree.red)):
    86:                        x = self.__rotRight()
    87:                        if (x != None):
    88:                            self.__copy(x)
    89:                        self.__color = RedBlackTree.black
Py  90:                        right = self.__right
    91:                        if (right != None):
    92:                           right.__color = RedBlackTree.red
    93:                           
    94:        else:
Py  95:            self.__RBinsertRight(val,1)
    96:            
    97:            # on the way back up check if a rotation is needed
    98:            # if search path has two red links with same orientation
    99:            # do a single rotation and flip the color bits
Py 100:            right = self.__right
   101:            if ((self.__color == RedBlackTree.red)and(right != None)and(right.__color == RedBlackTree.red)and(sw == 0)):
   102:                x = self.__rotLeft()
   103:                if (x != None):
   104:                    self.__copy(x)
Py 105:           
   106:           # flip the color bits
   107:            right = self.__right
   108:            if (right != None):
   109:                rr = right.__right
Py 110:                if (rr != None):
   111:                    if ((right.__color == RedBlackTree.red)and(rr.__color == RedBlackTree.red)):
   112:                        x = self.__rotLeft()
   113:                        if (x != None):
   114:                            self.__copy(x)
Py 115:                        self.__color = RedBlackTree.black
   116:                        left = self.__left
   117:                        if (left != None):
   118:                           left.__color = RedBlackTree.red
   119:
Py 120:
   121:# ============================================================
   122:# public methods
   123:# ============================================================
   124:    # constructor
Py 125:    def __init__(self,val = None):
   126:        self.__left      = None
   127:        self.__right     = None
   128:        self.__val       = val
   129:        self.__color     = RedBlackTree.red
Py 130:
   131:    # to string
   132:    def __str__(self):
   133:        return "[" + str(self.__val) + "," + str(self.__color) + "]"
   134:    
Py 135:    # accessor
   136:    def val(self):
   137:        return self.__val
   138:        
   139:    # accessor
Py 140:    def color(self):
   141:        return self.__color        
   142:
   143:    # recursive 'find'
   144:    def find(self,key):
Py 145:        result = None
   146:        if (key == self.__val):
   147:            result = self
   148:        elif (key < self.__val):
   149:            if (self.__left != None):
Py 150:                result = self.__left.find(key)
   151:        else:
   152:            if (self.__right != None):
   153:                result = self.__right.find(key)
   154:        return result
Py 155:
   156:    # inorder traversal using a class as the visitor
   157:    def inorderClass(self,visitor,depth):
   158:        if self.__left != None:
   159:            self.__left.inorderClass(visitor,depth+1)
Py 160:        visitor.visit(self,depth)
   161:        if self.__right != None:
   162:            self.__right.inorderClass(visitor,depth+1)
   163:            
   164:
Py 165:    # inorder traversal using a function as the visitor
   166:    def inorderFunction(self,visit,depth):
   167:        if self.__left != None:
   168:            self.__left.inorderFunction(visit,depth+1)
   169:        visit(self,depth)
Py 170:        if self.__right != None:
   171:            self.__right.inorderFunction(visit,depth+1)
   172:            
   173:    # main public 'insert' function        
   174:    def insert(self,val):
Py 175:        self.__RBinsert(val,0);
   176:        self.__color = RedBlackTree.black
   177:
   178:# define a visitor class if you need the visitor to have state     
   179:class NodeVisitor:
Py 180:    def __init__(self):
   181:        pass
   182:        
   183:    def visit(self,node,depth):
   184:        if (node.val() != None):
Py 185:            print "(" + str(node.val()) + ":"+ str(node.color()) + ":" + str(depth) + "),",
   186:            
   187:# just use a function if your visitor doesn't need state
   188:def visit(node,depth):
   189:    if (node.val() != None):
Py 190:        print "(" + str(node.val()) + ":"+ str(node.color()) + ":" + str(depth) + "),",
   191:    
   192:# ==================================
   193:# test program
   194:# ==================================
Py 195:if __name__ == "__main__":
   196:    nodelist = [11,4,8,14,17,6,9,7,16,15,13,5,19,18,12,10,3,20,1]
   197:    root  = RedBlackTree(2)
   198:    for n in nodelist:
   199:        root.insert(n)
Py 200:        
   201:    # pass a class instance
   202:    v = NodeVisitor()
   203:    print "Python C =",    
   204:    root.inorderClass(v,0)
Py 205:    print
   206:
   207:    # use a named function
   208:    print "Python F =",    
   209:    root.inorderFunction(visit,0)
Py 210:    print
   211:    
   212:    # do a find
   213:    t = root.find(16)
   214:    print str(t)
Py 215:
   216:
   217:
   218:
   219:
Py 220:
   221: