''' mbinary ######################################################################### # File : redBlackTree.py # Author: mbinary # Mail: zhuheqin1@gmail.com # Blog: https://mbinary.coding.me # Github: https://github.com/mbinary # Created Time: 2018-07-14 16:15 # Description: ######################################################################### ''' from functools import total_ordering from random import randint, shuffle @total_ordering class node: def __init__(self,val,left=None,right=None,isBlack=False): self.val =val self.left = left self.right = right self.parent= None self.isBlack = isBlack def __lt__(self,nd): return self.val < nd.val def __eq__(self,nd): return nd is not None and self.val == nd.val def setChild(self,nd,isLeft): if isLeft: self.left = nd else: self.right = nd if nd is not None: nd.parent = self def getChild(self,isLeft): if isLeft: return self.left else: return self.right def __bool__(self): return self.val is not None def __str__(self): color = 'B' if self.isBlack else 'R' val = '-' if self.parent==None else self.parent.val return f'{color}-{self.val}' def __repr__(self): return f'node({self.val},isBlack={self.isBlack})' class redBlackTree: def __init__(self,unique=False): '''if unique is True, all node'vals are unique, else there may be equal vals''' self.root = None self.unique = unique @staticmethod def checkBlack(nd): return nd is None or nd.isBlack @staticmethod def setBlack(nd,isBlack): if nd is not None: if isBlack is None or isBlack: nd.isBlack = True else:nd.isBlack = False def setRoot(self,nd): if nd is not None: nd.parent=None self.root= nd def find(self,val): nd = self.root while nd: if nd.val ==val: return nd else: nd = nd.getChild(nd.val>val) def getSuccessor(self,nd): if nd: if nd.right: nd = nd.right while nd.left: nd = nd.left return nd else: while nd.parent is not None and nd.parent.right is nd: nd = nd.parent return None if nd is self.root else nd.parent def rotate(self,prt,chd): '''rotate prt with the center of chd''' if self.root is prt: self.setRoot(chd) else: prt.parent.setChild(chd, prt.parent.left is prt) isLeftChd = prt.left is chd prt.setChild(chd.getChild(not isLeftChd), isLeftChd) chd.setChild(prt,not isLeftChd) def insert(self,nd): if nd.isBlack: nd.isBlack = False if self.root is None: self.setRoot(nd) self.root.isBlack = True else: parent = self.root while parent: if parent == nd : return None isLeft = parent > nd chd = parent.getChild(isLeft) if chd is None: parent.setChild(nd,isLeft) break else: parent = chd self.fixUpInsert(parent,nd) def fixUpInsert(self,parent,nd): ''' adjust color and level, there are two red nodes: the new one and its parent''' while not self.checkBlack(parent): grand = parent.parent isLeftPrt = grand.left is parent uncle = grand.getChild(not isLeftPrt) if not self.checkBlack(uncle): # case 1: new node's uncle is red self.setBlack(grand, False) self.setBlack(grand.left, True) self.setBlack(grand.right, True) nd = grand parent = nd.parent else: # case 2: new node's uncle is black(including nil leaf) isLeftNode = parent.left is nd if isLeftNode ^ isLeftPrt: # case 2.1 the new node is inserted in left-right or right-left form # grand grand # parent or parent # nd nd self.rotate(parent,nd) #parent rotate nd,parent = parent,nd # case 3 (case 2.2) the new node is inserted in left-left or right-right form # grand grand # parent or parent # nd nd self.setBlack(grand, False) self.setBlack(parent, True) self.rotate(grand,parent) self.setBlack(self.root,True) def copyNode(self,src,des): '''when deleting a node which has two kids, copy its succesor's data to his position data exclude left, right , isBlack ''' des.val = src.val def delete(self,val): '''delete node in a binary search tree''' if isinstance(val,node): val = val.val nd = self.find(val) if nd is None: return self._delete(nd) def _delete(self,nd): y = None if nd.left and nd.right: y= self.getSuccessor(nd) else: y = nd py = y.parent x = y.left if y.left else y.right if py is None: self.setRoot(x) else: py.setChild(x,py.left is y) if y != nd: self.copyNode(y,nd) if self.checkBlack(y): self.fixUpDel(py,x) def fixUpDel(self,prt,chd): ''' adjust colors and rotate ''' while self.root != chd and self.checkBlack(chd): isLeft =prt.left is chd brother = prt.getChild(not isLeft) # brother is black lb = self.checkBlack(brother.getChild(isLeft)) rb = self.checkBlack(brother.getChild(not isLeft)) if not self.checkBlack(brother): # case 1: brother is red. converted to case 2,3,4 self.setBlack(prt,False) self.setBlack(brother,True) self.rotate(prt,brother) elif lb and rb: # case 2: brother is black and two kids are black. # conveted to the begin case self.setBlack(brother,False) chd = prt prt= chd.parent else: if rb: # case 3: brother is black and left kid is red and right child is black # rotate bro to make g w wl wr in one line # uncle's son is nephew, and niece for uncle's daughter nephew = brother.getChild(isLeft) self.setBlack(nephew,True) self.setBlack(brother,False) # brother (not isLeft) rotate self.rotate(brother,nephew) brother = nephew # case 4: brother is black and right child is red brother.isBlack = prt.isBlack self.setBlack(prt,True) self.setBlack(brother.getChild(not isLeft),True) self.rotate(prt,brother) chd = self.root self.setBlack(chd,True) def sort(self,reverse = False): ''' return a generator of sorted data''' def inOrder(root): if root is None:return if reverse: yield from inOrder(root.right) else: yield from inOrder(root.left) yield root if reverse: yield from inOrder(root.left) else: yield from inOrder(root.right) yield from inOrder(self.root) def display(self): def getHeight(nd): if nd is None:return 0 return max(getHeight(nd.left),getHeight(nd.right)) +1 def levelVisit(root): from collections import deque lst = deque([root]) level = [] h = getHeight(root) ct = lv = 0 while 1: ct+=1 nd = lst.popleft() if ct >= 2**lv: lv+=1 if lv>h:break level.append([]) level[-1].append(str(nd)) if nd is not None: lst += [nd.left,nd.right] else: lst +=[None,None] return level def addBlank(lines): width = 1+len(str(self.root)) sep = ' '*width n = len(lines) for i,oneline in enumerate(lines): k = 2**(n-i) -1 new = [sep*((k-1)//2)] for s in oneline: new.append(s.ljust(width)) new.append(sep*k) lines[i] = new return lines lines = levelVisit(self.root) lines = addBlank(lines) li = [''.join(line) for line in lines] length = 10 if li==[] else max(len(i) for i in li)//2 begin ='\n'+ 'red-black-tree'.rjust(length+14,'-') + '-'*(length) end = '-'*(length*2+14)+'\n' return '\n'.join([begin,*li,end]) def __str__(self): return self.display() def genNum(n =10): nums =[] for i in range(n): while 1: d = randint(0,100) if d not in nums: nums.append(d) break return nums def buildTree(n=10,nums=None,visitor=None): if nums is None or nums ==[]: nums = genNum(n) rbtree = redBlackTree() print(f'build a red-black tree using {nums}') for i in nums: rbtree.insert(node(i)) if visitor: visitor(rbtree,i) return rbtree,nums def testInsert(nums=None): def visitor(t,val): print('inserting', val) print(t) rbtree,nums = buildTree(visitor = visitor,nums=nums) print('-'*5+ 'in-order visit' + '-'*5) for i,j in enumerate(rbtree.sort()): print(f'{i+1}: {j}') def testSuc(nums=None): rbtree,nums = buildTree(nums=nums) for i in rbtree.sort(): print(f'{i}\'s suc is {rbtree.getSuccessor(i)}') def testDelete(nums=None): rbtree,nums = buildTree(nums = nums) print(rbtree) for i in sorted(nums): print(f'deleting {i}') rbtree.delete(i) print(rbtree) if __name__=='__main__': lst =[45, 30, 64, 36, 95, 38, 76, 34, 50, 1] lst = [0,3,5,6,26,25,8,19,15,16,17] #testSuc(lst) #testInsert(lst) testDelete()