1 #!/usr/bin/env python  2 #coding: gbk  3   4 ########################################################################  5 #Author: Feng Ruohang  6 #Create: 2014/10/06 11:38  7 #Digest: Provide a common data struct: Red Black Tree  8 ########################################################################  9  10 class RBT(object): 11     class Node(object): 12         ‘‘‘ 13         Node used in RBTree 14         ‘‘‘ 15         def __init__(self,key,value=http://www.mamicode.com/None,color=False,N=1): 16             self.key = key 17             self.val = value 18             self.color = color   #False for Black, True for Red 19             self.N = N           #Total numbers of nodes in this subtree 20             self.left = None 21             self.right = None 22  23         def __cmp__(l,r): 24             return cmp(l.key,r.key) 25  26         def __eq__(l,r): 27             return True if l.key == r.key else False 28  29         def __add__(l,r): 30             l.value + r.value 31  32         def reduce(self,new_val): 33             self.val.extend(new_val) 34  35     def __init__(self): 36         self.root = None 37  38     #====================APIs====================# 39     #=====Basic API 40  41     def get(self,key): 42         return  self.__get(self.root,key) 43  44     def put(self,key,val): 45         self.root = self.__put(self.root,key,val) 46         self.root.color = False  47  48     def append(self,key,val): 49         self.root = self.__append(self.root,key,val) 50         self.root.color = False 51  52     def delete(self,key): 53         if not self.contains(key): 54             raise LookupError(‘No such keys in rbtree. Fail to Delete‘) 55         if self.__is_black(self.root.left) and self.__is_black(self.root.right): 56             self.root.color = True 57         self.root = self.__delete(self.root,key) 58         if not self.is_empty():  59             self.root.color = False 60  61     def del_min(self): 62         if self.is_empty():  63             raise LookupError(‘Empty Red-Black Tree. Can\‘t delete min‘) 64         if self.__is_black(self.root.left) and self.__is_black(self.root.right): 65             self.root.color = True 66         self.root = self.__del_min(self.root) 67         if not self.is_empty(): self.root.color = False 68  69     def del_max(self): 70         if self.is_empty(): 71             raise LookupError(‘Empty Red-Black Tree. Can\‘t delete max‘) 72         if self.__is_black(self.root.left) and self.__is_black(self.root.right): 73             self.root.color = True 74         self.root = self.__del_max(self.root) 75         if not self.is_empty(): self.root.color = False  76  77     def size(self): 78         return self.__size(self.root) 79  80     def is_empty(self): 81         return not self.root 82  83     def contains(self,key): 84         return bool(self.get(key)) 85  86     #=====Advance API 87     def min(self): 88         if self.is_empty(): 89             return None 90         return self.__min(self.root).key 91  92     def max(self): 93         if self.is_empty(): 94             return None 95         return self.__max(self.root).key 96  97     def floor(self,key): 98         x = self.__floor(self.root,key) 99         if x:100             return x.key,x.val101         else:102             return None,None103 104     def ceil(self,key):105         x = self.__ceil(self.root,key)106         if x:107             return x.key,x.val108         else:109             return None,None110 111     def below(self,key):112         index = self.index(key)113         if not 0 <= index - 1 < self.size():114             return None,None    #Return None if out of range115         x = self.__select(self.root,index - 1)116         return x.key,x.val117 118     def above(self,key):119         index = self.index(key)120         if self.contains(key):121             if not 0 <= index + 1 < self.size():122                 return None,None    #Return None if out of range123             else:124                 x = self.__select(self.root,index+1)125                 return x.key,x.val126         else:#if key is not in tree. then select(i) is what we need127             if not 0 <= index < self.size():128                 return None,None    #Return None if out of range129             else:130                 x = self.__select(self.root,index)131                 return x.key,x.val132 133     def index(self,key):134         return self.__index(self.root,key) 135 136     def keys(self):137         ‘‘‘Return All Keys in the tree ‘‘‘138         return self.range(self.min(),self.max())139 140     def range(self,lo,hi):141         ‘‘‘Take two keys. return keys between them‘‘‘142         q = []143         self.__range(self.root,q,lo,hi)144         return q145 146     def select(self,index):147         ‘‘‘Given Index Return Corresponding key ‘‘‘148         if not 0 <= index < self.size():149             return None150         return self.__select(self.root,index).key151 152     def width(self,lo,hi):153         ‘‘‘Return the numbers of keys between lo and hi ‘‘‘154         if lo > hi: 155             return 0156         if self.contains(hi):157             return self.index(hi) - self.index(lo) + 1158         else:159             return self.index(hi) - self.index(lo)160 161 162     #===============Private Method===============#163     #=====Basic164     def __get(self,x,key):165         while x:166             tag = cmp(key,x.key)167             if tag < 0 : x = x.left168             elif tag > 0 :x = x.right169             else: return x.val170 171     def __put(self,h,key,val):172         if not h: 173             return self.Node(key,val,True,1)174         tag = cmp(key,h.key)175         if tag < 0: 176             h.left = self.__put(h.left,key,val)177         elif tag > 0: 178             h.right = self.__put(h.right,key,val)179         else: 180             h.val = val   #Update181 182         if self.__is_black(h.left) and self.__is_red(h.right):183             h = self.__rotate_left(h)184         if self.__is_red(h.left) and self.__is_red(h.left.left):185             h = self.__rotate_right(h)186         if self.__is_red(h.left) and self.__is_red(h.right):187             self.__flip_colors(h)188         h.N = self.__size(h.left) + self.__size(h.right) + 1189         return h190 191     def __append(self,h,key,val):192         if not h: 193             return self.Node(key,val,True,1)194         tag = cmp(key,h.key)195         if tag < 0: 196             h.left = self.__append(h.left,key,val)197         elif tag > 0: 198             h.right = self.__append(h.right,key,val)199         else: 200             h.reduce(val)   #append.201 202         if self.__is_black(h.left) and self.__is_red(h.right):203             h = self.__rotate_left(h)204         if self.__is_red(h.left) and self.__is_red(h.left.left):205             h = self.__rotate_right(h)206         if self.__is_red(h.left) and self.__is_red(h.right):207             self.__flip_colors(h)208         h.N = self.__size(h.left) + self.__size(h.right) + 1209         return h210 211     def __del_min(self,h):212         if not h.left: #if h is empty:return None213             return None214 215         if self.__is_black(h.left) and self.__is_black(h.left.left):216             self.__move_red_left(h)217         h.left = self.__del_min(h.left) #Del recursive218         return self.__balance(h)219 220     def __del_max(self,h):221         if self.__is_red(h.left): 222             h = self.__rotate_right(h)223         if not h.right: 224             return None225         if self.__is_black(h.right) and self.__is_black(h.right.left):226             h = self.__move_red_right(h)227         h.right = self.__del_max(h.right)228         return self.__balance(h)229 230     def __delete(self,h,key):231         if key < h.key:232             if self.__is_black(h.left) and self.__is_black(h.left.left):233                 h = self.__move_red_left(h)234             h.left = self.__delete(h.left,key)235         else:236             if self.__is_red(h.left):237                 h = self.__rotate_right(h)238             if key == h.key and not h.right:239                 return None240             if self.__is_black(h.right) and self.__is_black(h.right.left):241                 h = self.__move_red_right(h)242             if key == h.key:#replace h with min of right subtree243                 x = self.__min(h.right)244                 h.key = x.key245                 h.val = x.val246                 h.right = self.__del_min(h.right)247             else:248                 h.right = self.__delete(h.right,key)249         h = self.__balance(h)250         return h251     252     #=====Advance253     def __min(self,h):254         #Assume h is not null255         if not h.left:256             return h257         else:258             return self.__min(h.left)259 260     def __max(self,h):261         #Assume h is not null262         if not h.right:263             return h264         else:265             return self.__max(h.right)266 267     def __floor(self,h,key):268         ‘‘‘Find the NODE with key <= given key in the tree rooted at h ‘‘‘269         if not h:270             return None271         tag = cmp(key,h.key)272         if tag == 0:273             return h274         if tag < 0:275             return self.__floor(h.left,key)276         t = self.__floor(h.right,key)277         if t:#if find in right tree278             return t279         else:#else return itself280             return h281 282     def __ceil(self,h,key):283         ‘‘‘Find the NODE with key >= given key in the tree rooted at h ‘‘‘284         if not h:285             return None286         tag = cmp(key,h.key)287         if tag == 0:288             return h289         if tag > 0: # key is bigger290             return self.__ceil(h.right,key)291         t = self.__ceil(h.left,key)#key is lower.Try to find ceil left292         if t:#if find in left tree293             return t294         else:#else return itself295             return h296 297     def __index(self,h,key):298         if not h:299             return 0300         tag = cmp(key,h.key)301         if tag < 0:302             return self.__index(h.left,key)303         elif tag > 0:   #Key is bigger304             return self.__index(h.right,key) + 1 + self.__size(h.left)305         else:   #Eq306             return self.__size(h.left)307 308     def __select(self,h,index):309         ‘‘‘assert h. assert 0 <= index < size(tree) ‘‘‘310         l_size = self.__size(h.left)311         if l_size > index:312             return self.__select(h.left,index)313         elif l_size < index:314             return self.__select(h.right,index - l_size - 1)315         else:316             return h317 318     def __range(self,h,q,lo,hi):319         if not h: 320             return321         tag_lo = cmp(lo,h.key)322         tag_hi = cmp(hi,h.key)323         if tag_lo < 0:#lo key is lower than h.key324             self.__range(h.left,q,lo,hi)325         if tag_lo <= 0 and tag_hi >= 0:326             q.append(h.key)327         if tag_hi > 0 :# hi key is bigger than h.key328             self.__range(h.right,q,lo,hi)329 330 331     #===============Adjust Functions=============#332     def __rotate_right(self,h):333         x = h.left334         h.left,x.right = x.right,h335         x.color,x.N = h.color,h.N336         h.color,h.N = True,self.__size(h.left) + self.__size(h.right) + 1337         return x338 339     def __rotate_left(self,h):340         x = h.right341         h.right,x.left = x.left,h342         x.color,x.N = h.color,h.N343         h.color,h.N = True,self.__size(h.left) + self.__size(h.right) + 1344         return x345 346     def __flip_colors(self,h):347         h.color = not h.color348         h.left.color = not h.left.color349         h.right.color = not h.right.color350 351     def __move_red_left(self,h):352         self.__flip_colors(h)353         if self.__is_red(h.right.left):354             h = self.__rotate_left(h)355         return h356 357     def __move_red_right(self,h):358         self.__flip_colors(h)359         if self.__is_red(h.left.left):360             h = self.__rotate_right(h)361         return h362 363     def __balance(self,h):364         if self.__is_red(h.right): 365             h = self.__rotate_left(h)366         if self.__is_red(h.left) and self.__is_red(h.left.left): 367             h = self.__rotate_right(h)368         if self.__is_red(h.left) and self.__is_red(h.right):369             self.__flip_colors(h)370         h.N = self.__size(h.left) + self.__size(h.right) + 1371         return h372     373     #Class Method374     @staticmethod375     def __is_red(x):376         return False if not x else x.color377 378     @staticmethod379     def __is_black(x):380         return True  if not x else not x.color381 382     @staticmethod383     def __size(x):384         return 0 if not x else x.N385 386 387 def RBT_testing():388     ‘‘‘API Examples ‘‘‘389     t = RBT()390     test_data = http://www.mamicode.com/"SEARCHXMPL"391 392     print ‘=====testing is_empty()\nBefore Insertion‘393     print t.is_empty()394 395     for letter in test_data:396         t.put(letter,[ord(letter)])397         print "Test Inserting:%s, tree size is %d" % (letter,t.size())398     print "After insertion it return:"399     print t.is_empty()400     print  "====test is_empty complete\n"401 402 403     print "=====Tesing Get method:"404     print "get ‘s‘ is "405     print t.get(‘S‘)406     print "get ‘H‘ is "407     print t.get(‘H‘)408 409     print ‘==Trying get null key: get "F" is‘410     print t.get(‘F‘)411     print "=====Testing Get method end\n\n"412 413     print "=====Testing ceil and floor"414     print "Ceil(‘L‘)"415     print t.ceil(‘L‘)416     print "Ceil(‘F‘) *F is not in tree"417     print t.ceil(‘F‘)418 419     print "Floor(‘L‘)"420     print t.ceil(‘L‘)421     print "Floor(‘F‘)"422     print t.ceil(‘F‘)423 424     print ‘======test append method‘425     print ‘Orient key e is correspond with‘426     print t.get(‘E‘)427     t.append(‘E‘,[4])428     print ‘==After append‘429     print t.get(‘E‘)430     print "=====Testing Append method end\n\n" 431 432     print "=====Testing index()"433     print "index(E)" 434     print t.index(‘E‘)435     print "index(L),select(4)"436     print t.index(‘L‘),t.select(4)437     print "index(‘M‘),select(5)"438     print t.index(‘M‘),t.select(5)439     print "index a key not in tree:\n index(‘N‘),select(6)"440     print t.index(‘N‘),t.select(6)441     print "index(‘P‘)"442     print t.index(‘P‘)443     444     445     print "=====Testing select"446     print "select(3) = "447     print t.select(3)448     print "select and index end...\n\n"449 450     print "====Tesing Min and Max"451     print "min key is:"452     print t.min()453     print "max key is"454     print t.max()455 456     print "==How much between min and max:"457     print t.width(t.min(),t.max())458     print "keys between min and max:"459     print t.keys()460     print "keys in ‘E‘ and ‘M‘ "461     print t.range(‘E‘,‘M‘)462 463 464     print "try to delete min_key:"465     print "But we could try contains(‘A‘) first"466     print t.contains(‘A‘)467     t.del_min()468     print "After deletion t.contains(‘A‘) is "469     print t.contains(‘A‘)470 471     print t.min()472     print "try to kill one more min key:"473     t.del_min()474     print t.min()475     print "try to delete max_key,New Max key is :"476     t.del_max()477     print t.max()478     print "=====Tesing Min and Max complete\n\n"479 480 481 482     print ‘=====Deleting Test‘483     print t.size()484     t.delete(‘H‘)485     print t.size()486 487     print ‘Delete a non-exists key:‘488     try:    489         t.delete(‘F‘)490     except:491         print "*Look up error occur*"492 493     print "=====Testing Delete method complete"494 495 def test_basic_api():496     print "==========Testing Basic API=========="497     t = RBT()498     print "Test Data: FENGDIJKABCLM"499     test_data = http://www.mamicode.com/"FENGDIJKABCLM" #from A-N,without H500 501     #=====put()502     print "==========put() test begin!=========="503     for letter in test_data:504         t.put(letter,[ord(letter)]) #Value is [ascii order of letter]505         print "put(%s); Now tree size is %d"%(letter,t.size())506     print ‘Final tree size is %d‘%t.size()507     print "==========put() test complete!==========\n"508 509     #=====get()510     print "==========get() test begin!=========="511     print "get(‘F‘):\t%s"%repr(t.get(‘F‘))512     print "get(‘A‘):\t%s"%repr(t.get(‘A‘))513     print "get a non-exist key Z: get(‘Z‘):\t%s"%repr(t.get(‘Z‘))514     print "==========get() test complete!==========\n"515 516     #=====append()517     print "=====append() test begin!=========="518     print "First append to a exist key:[F]"519     print "Before Append:get(‘F‘):\t%s"%repr(t.get(‘F‘))520     print "append(‘F‘,[3,‘haha‘]):\t%s"%repr(t.append(‘F‘,[3,‘haha‘]))521     print "After Append:get(‘F‘):\t%s\n"%repr(t.get(‘F‘))522     print "Second append to a non-exist key:[O]"523     print "Before Append:get(‘O‘):\t%s"%repr(t.get(‘O‘))524     print "append a non-exist key O: append(‘O‘,[‘value of O‘]):\t%s"%repr(t.append(‘O‘,[‘value of O‘]))525     print "After Append:get(‘O‘):\t%s\n"%repr(t.get(‘O‘))526     print "==========append() test complete!==========\n"527 528     #=====delete()529     print "==========delete() test begin!=========="530     test_data2 = [x for x in test_data]531     test_data2.reverse()532     for letter in test_data2:533         t.delete(letter)534         print "delete(%s); Now tree size is %d"%(letter,t.size())535     print ‘Final tree size is %d‘%t.size()536     print "==========delete() test complete!==========\n"537 538     print "==========Basic API Test Complete==========\n\n"539 540 def test_advance_api():541     print "==========Testing min max floor ceil above below =========="542     t = RBT()543     print "Test Data: FENGDIJKABCLM"544     test_data = http://www.mamicode.com/"FENGDIJKABCLM" #from A-N,without H545     for letter in test_data:546         t.put(letter,[ord(letter)]) #Value is [ascii order of letter]547 548     #=====min() and del_min()549     print "==========min() and del_min() test begin!=========="550     print "Original min():\t%s"%repr(t.min())551     print "run del_min()"552     t.del_min()553     print "After del_min:min()\t%s"%repr(t.min())554 555     print "run del_min() again"556     t.del_min()557     print "After del_min run again:min()\t%s"%repr(t.min())558 559     print "=====max() and del_max() test begin!"560     print "Original max():\t%s"%repr(t.max())561     print "run del_max()"562     t.del_max()563     print "After del_max:max()\t%s"%repr(t.max())564 565     print "run del_max() again"566     t.del_max()567     print "After del_max run again:max()\t%s"%repr(t.max())568     print "==========min() max() del_min() del_max() test complete!==========\n"569 570 def test_int_api():571     #======ceil floor above below572     print "==========Testing ceil floor above below =========="573     t = RBT()574     print "Test Data: FENGDIJKABCLM - [AHN] = FEGDIJKBCLM"575     test_data = http://www.mamicode.com/"FEGDIJKBCLM" #from A-N, Del A H N576 577     for letter in test_data:578         t.put(letter,[ord(letter)]) #Value is [ascii order of letter]579     print "Node\tceil\t\tfloor\t\tabove\t\tbelow"580     for P in [‘A‘,‘B‘,‘C‘,‘G‘,‘H‘,‘I‘,‘L‘,‘M‘,‘N‘]:581         print "%s\t%s\t%s\t%s\t%s"%(P,t.ceil(P),t.floor(P),t.above(P),t.below(P))582 583 if __name__ == ‘__main__‘:584     test_basic_api()585     test_advance_api()586     test_int_api()