红黑树的Python实现

想用红黑树，怎么搜都搜不到现成的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树与红黑树的相同之处，才能对那些乱七八糟的插入删除调整操作有直观的认识。