python 红黑树_红黑树的Python实现

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 classRBT(object):11 classNode(object):12 '''

13 Node used in RBTree14 '''

15 def __init__(self,key,value=None,color=False,N=1):16 self.key =key17 self.val =value18 self.color = color #False for Black, True for Red

19 self.N = N #Total numbers of nodes in this subtree

20 self.left =None21 self.right =None22

23 def __cmp__(l,r):24 returncmp(l.key,r.key)25

26 def __eq__(l,r):27 return True if l.key == r.key elseFalse28

29 def __add__(l,r):30 l.value +r.value31

32 defreduce(self,new_val):33 self.val.extend(new_val)34

35 def __init__(self):36 self.root =None37

38 #====================APIs====================#

39 #=====Basic API

40

41 defget(self,key):42 return self.__get(self.root,key)43

44 defput(self,key,val):45 self.root = self.__put(self.root,key,val)46 self.root.color =False47

48 defappend(self,key,val):49 self.root = self.__append(self.root,key,val)50 self.root.color =False51

52 defdelete(self,key):53 if notself.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 =True57 self.root = self.__delete(self.root,key)58 if notself.is_empty():59 self.root.color =False60

61 defdel_min(self):62 ifself.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 =True66 self.root = self.__del_min(self.root)67 if not self.is_empty(): self.root.color =False68

69 defdel_max(self):70 ifself.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 =True74 self.root = self.__del_max(self.root)75 if not self.is_empty(): self.root.color =False76

77 defsize(self):78 return self.__size(self.root)79

80 defis_empty(self):81 return notself.root82

83 defcontains(self,key):84 returnbool(self.get(key))85

86 #=====Advance API

87 defmin(self):88 ifself.is_empty():89 returnNone90 return self.__min(self.root).key91

92 defmax(self):93 ifself.is_empty():94 returnNone95 return self.__max(self.root).key96

97 deffloor(self,key):98 x = self.__floor(self.root,key)99 ifx:100 returnx.key,x.val101 else:102 returnNone,None103

104 defceil(self,key):105 x = self.__ceil(self.root,key)106 ifx:107 returnx.key,x.val108 else:109 returnNone,None110

111 defbelow(self,key):112 index =self.index(key)113 if not 0 <= index - 1

115 x = self.__select(self.root,index - 1)116 returnx.key,x.val117

118 defabove(self,key):119 index =self.index(key)120 ifself.contains(key):121 if not 0 <= index + 1

123 else:124 x = self.__select(self.root,index+1)125 returnx.key,x.val126 else:#if key is not in tree. then select(i) is what we need

127 if not 0 <= index

129 else:130 x = self.__select(self.root,index)131 returnx.key,x.val132

133 defindex(self,key):134 return self.__index(self.root,key)135

136 defkeys(self):137 '''Return All Keys in the tree'''

138 returnself.range(self.min(),self.max())139

140 defrange(self,lo,hi):141 '''Take two keys. return keys between them'''

142 q =[]143 self.__range(self.root,q,lo,hi)144 returnq145

146 defselect(self,index):147 '''Given Index Return Corresponding key'''

148 if not 0 <= index

152 defwidth(self,lo,hi):153 '''Return the numbers of keys between lo and hi'''

154 if lo >hi:155 return0156 ifself.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 whilex:166 tag =cmp(key,x.key)167 if tag < 0 : x =x.left168 elif tag > 0 :x =x.right169 else: returnx.val170

171 def __put(self,h,key,val):172 if noth: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 returnh190

191 def __append(self,h,key,val):192 if noth: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 returnh210

211 def __del_min(self,h):212 if not h.left: #if h is empty:return None

213 returnNone214

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 noth.right:224 returnNone225 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 <232 if self.__is_black and h="self.__move_red_left(h)234" h.left="self.__delete(h.left,key)235" else:236 self.__is_red key="=" h.key noth.right:239 returnnone240 h.key: with min of right subtree>

243 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 returnh251

252 #=====Advance

253 def __min(self,h):254 #Assume h is not null

255 if noth.left:256 returnh257 else:258 return self.__min(h.left)259

260 def __max(self,h):261 #Assume h is not null

262 if noth.right:263 returnh264 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 noth:270 returnNone271 tag =cmp(key,h.key)272 if tag ==0:273 returnh274 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 returnt279 else:#else return itself

280 returnh281

282 def __ceil(self,h,key):283 '''Find the NODE with key >= given key in the tree rooted at h'''

284 if noth:285 returnNone286 tag =cmp(key,h.key)287 if tag ==0:288 returnh289 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 returnt294 else:#else return itself

295 returnh296

297 def __index(self,h,key):298 if noth:299 return0300 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 <314 return self.__select l_size else:316 returnh317>

318 def __range(self,h,q,lo,hi):319 if noth: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.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) + 1

337 returnx338

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) + 1

344 returnx345

346 def __flip_colors(self,h):347 h.color = noth.color348 h.left.color = noth.left.color349 h.right.color = noth.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 returnh356

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 returnh362

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 returnh372

373 #Class Method

374 @staticmethod375 def __is_red(x):376 return False if not x elsex.color377

378 @staticmethod379 def __is_black(x):380 return True if not x else notx.color381

382 @staticmethod383 def __size(x):384 return 0 if not x elsex.N385

386

387 defRBT_testing():388 '''API Examples'''

389 t =RBT()390 test_data = "SEARCHXMPL"

391

392 print '=====testing is_empty()\nBefore Insertion'

393 printt.is_empty()394

395 for letter intest_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 printt.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 printt.min()453 print "max key is"

454 printt.max()455

456 print "==How much between min and max:"

457 printt.width(t.min(),t.max())458 print "keys between min and max:"

459 printt.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 printt.min()472 print "try to kill one more min key:"

473 t.del_min()474 printt.min()475 print "try to delete max_key,New Max key is :"

476 t.del_max()477 printt.max()478 print "=====Tesing Min and Max complete\n\n"

479

480

481

482 print '=====Deleting Test'

483 printt.size()484 t.delete('H')485 printt.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 deftest_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 intest_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 intest_data]531 test_data2.reverse()532 for letter intest_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 deftest_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 intest_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 deftest_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 intest_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()

314>232>