-
Notifications
You must be signed in to change notification settings - Fork 88
/
redBlackTree0.py
339 lines (311 loc) · 11.1 KB
/
redBlackTree0.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
''' mbinary
#########################################################################
# File : redBlackTree.py
# Author: mbinary
# Mail: zhuheqin1@gmail.com
# Blog: https://mbinary.xyz
# 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.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
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'
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 getParent(self,chd):
'''note that use is to find real node when different nodes have euqiv val'''
if self.root is chd:return None
nd = self.root
while nd:
if nd>chd and nd.left is not None:
if nd.left is chd: return nd
else: nd = nd.left
elif nd<chd and nd.right is not None:
if nd.right is chd: return nd
else: nd = nd.right
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:return self.getParent(nd)
def setRoot(self,nd):
if nd is not None: nd.parent=None
self.root= nd
def transferParent(self,origin,new):
if origin is self.root:
self.root = new
else:
prt = self.getParent(origin)
prt.setChild(new, prt.left is origin)
def rotate(self,prt,chd,direction):
'''rotate prt, and return new prt, namyly the original chd'''
prt.setChild(chd.getChild(direction), not direction)
chd.setChild(prt,direction)
return chd
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 = self.getParent(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 = self.getParent(nd)
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,isLeftPrt) #parent rotate
grand.setChild(nd,isLeftPrt)
nd,parent = parent,nd
# case 2.2 the new node is inserted in left-left or right-right form
# grand grand
# parent or parent
# nd nd
# grand rotate
self.rotate(grand,parent,not isLeftPrt)
self.setBlack(grand, False)
self.setBlack(parent, True)
self.transferParent(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
y = None
if nd.left and nd.right:
y= self.getSuccessor(nd)
else:
y = nd
py = self.getParent(y)
x = y.left if y.left else y.right
if py is None:
self.root = 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
# prt (isLeft) rotate
self.rotate(prt,brother,isLeft)
self.setBlack(prt,False)
self.setBlack(brother,True)
self.transferParent(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 = self.getParent(chd)
else:
if rb:
# case 3: brother is black and left kid is red and right child is black
# 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
prt.setChild(nephew,not isLeft)
self.rotate(brother,nephew,not isLeft)
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)
# prt left rotate
self.rotate(prt,brother,isLeft)
self.transferParent(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 = 5
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]
li.insert(0,'red-black-tree'.rjust(48,'-') + '-'*33)
li.append('end'.rjust(42,'-')+'-'*39+'\n')
return '\n'.join(li)
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)
print(rbtree)
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 nums:
print(f'deleting {i}')
rbtree.delete(i)
print(rbtree)
if __name__=='__main__':
#lst = [41,38,31,12,19,8]
#lst.sort()
lst = None
lst =[45, 30, 64, 36, 95, 38, 76, 34, 50, 1]
testSuc()
testInsert(lst)
testDelete()