想用红黑树，怎么搜都搜不到现成的Python实现。干脆自己写一个。
算法的结构按照Sedgewick的《算法(4th)》一书第三章写成，略有改动。
完整的API实现，也用了一些测试case，暂时没发现问题。
这玩意就是好用，谁用谁知道。
废话不多说直接上代码。
#注意事项： 重载RBT.Node.reduce(self,new_val)来实现 append()方法中对已存在主键对应值的自定义合并操作。默认为调用list.extend()方法。
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=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.val 101 else: 102 return None,None 103 104 def ceil(self,key): 105 x = self.__ceil(self.root,key) 106 if x: 107 return x.key,x.val 108 else: 109 return None,None 110 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 range 115 x = self.__select(self.root,index - 1) 116 return x.key,x.val 117 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 range 123 else: 124 x = self.__select(self.root,index+1) 125 return x.key,x.val 126 else:#if key is not in tree. then select(i) is what we need 127 if not 0 <= index < self.size(): 128 return None,None #Return None if out of range 129 else: 130 x = self.__select(self.root,index) 131 return x.key,x.val 132 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 q 145 146 def select(self,index): 147 '''Given Index Return Corresponding key ''' 148 if not 0 <= index < self.size(): 149 return None 150 return self.__select(self.root,index).key 151 152 def width(self,lo,hi): 153 '''Return the numbers of keys between lo and hi ''' 154 if lo > hi: 155 return 0 156 if self.contains(hi): 157 return self.index(hi) - self.index(lo) + 1 158 else: 159 return self.index(hi) - self.index(lo) 160 161 162 #===============Private Method===============# 163 #=====Basic 164 def __get(self,x,key): 165 while x: 166 tag = cmp(key,x.key) 167 if tag < 0 : x = x.left 168 elif tag > 0 :x = x.right 169 else: return x.val 170 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 #Update 181 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) + 1 189 return h 190 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) + 1 209 return h 210 211 def __del_min(self,h): 212 if not h.left: #if h is empty:return None 213 return None 214 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 recursive 218 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 None 225 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 None 240 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 subtree 243 x = self.__min(h.right) 244 h.key = x.key 245 h.val = x.val 246 h.right = self.__del_min(h.right) 247 else: 248 h.right = self.__delete(h.right,key) 249 h = self.__balance(h) 250 return h 251 252 #=====Advance 253 def __min(self,h): 254 #Assume h is not null 255 if not h.left: 256 return h 257 else: 258 return self.__min(h.left) 259 260 def __max(self,h): 261 #Assume h is not null 262 if not h.right: 263 return h 264 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 None 271 tag = cmp(key,h.key) 272 if tag == 0: 273 return h 274 if tag < 0: 275 return self.__floor(h.left,key) 276 t = self.__floor(h.right,key) 277 if t:#if find in right tree 278 return t 279 else:#else return itself 280 return h 281 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 None 286 tag = cmp(key,h.key) 287 if tag == 0: 288 return h 289 if tag > 0: # key is bigger 290 return self.__ceil(h.right,key) 291 t = self.__ceil(h.left,key)#key is lower.Try to find ceil left 292 if t:#if find in left tree 293 return t 294 else:#else return itself 295 return h 296 297 def __index(self,h,key): 298 if not h: 299 return 0 300 tag = cmp(key,h.key) 301 if tag < 0: 302 return self.__index(h.left,key) 303 elif tag > 0: #Key is bigger 304 return self.__index(h.right,key) + 1 + self.__size(h.left) 305 else: #Eq 306 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 h 317 318 def __range(self,h,q,lo,hi): 319 if not h: 320 return 321 tag_lo = cmp(lo,h.key) 322 tag_hi = cmp(hi,h.key) 323 if tag_lo < 0:#lo key is lower than h.key 324 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.key 328 self.__range(h.right,q,lo,hi) 329 330 331 #===============Adjust Functions=============# 332 def __rotate_right(self,h): 333 x = h.left 334 h.left,x.right = x.right,h 335 x.color,x.N = h.color,h.N 336 h.color,h.N = True,self.__size(h.left) + self.__size(h.right) + 1 337 return x 338 339 def __rotate_left(self,h): 340 x = h.right 341 h.right,x.left = x.left,h 342 x.color,x.N = h.color,h.N 343 h.color,h.N = True,self.__size(h.left) + self.__size(h.right) + 1 344 return x 345 346 def __flip_colors(self,h): 347 h.color = not h.color 348 h.left.color = not h.left.color 349 h.right.color = not h.right.color 350 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 h 356 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 h 362 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) + 1 371 return h 372 373 #Class Method 374 @staticmethod 375 def __is_red(x): 376 return False if not x else x.color 377 378 @staticmethod 379 def __is_black(x): 380 return True if not x else not x.color 381 382 @staticmethod 383 def __size(x): 384 return 0 if not x else x.N 385 386 387 def RBT_testing(): 388 '''API Examples ''' 389 t = RBT() 390 test_data = "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 = "FENGDIJKABCLM" #from A-N,without H 500 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 = "FENGDIJKABCLM" #from A-N,without H 545 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 below 572 print "==========Testing ceil floor above below ==========" 573 t = RBT() 574 print "Test Data: FENGDIJKABCLM - [AHN] = FEGDIJKBCLM" 575 test_data = "FEGDIJKBCLM" #from A-N, Del A H N 576 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()
查找操作的数据结构不断进化，才有了红黑树：从链表到二叉平衡树，再到2-3树，最后到红黑树。
红黑树本质上是用二叉平衡树的形式来模拟2-3树的功能
《算法导论》也好，其他什么乱七八糟算法书博客也罢，讲红黑树都没讲到本质。
Sedgewick的《算法(4th)》这本书就很不错：起码他告诉你红黑树是怎么来的。
仔细理解2-3树与红黑树的相同之处，才能对那些乱七八糟的插入删除调整操作有直观的认识。
