forked from keon/algorithms
-
Notifications
You must be signed in to change notification settings - Fork 0
/
rbtree.py
294 lines (274 loc) · 11 KB
/
rbtree.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
"""
Implementation of Red-Black tree.
"""
class RBNode:
def __init__(self, val, is_red, parent=None, left=None, right=None):
self.val = val
self.parent = parent
self.left = left
self.right = right
self.color = is_red
class RBTree:
def __init__(self):
self.root = None
def left_rotate(self, node):
# set the node as the left child node of the current node's right node
right_node = node.right
if right_node is None:
return
else:
# right node's left node become the right node of current node
node.right = right_node.left
if right_node.left is not None:
right_node.left.parent = node
right_node.parent = node.parent
# check the parent case
if node.parent is None:
self.root = right_node
elif node is node.parent.left:
node.parent.left = right_node
else:
node.parent.right = right_node
right_node.left = node
node.parent = right_node
def right_rotate(self, node):
# set the node as the right child node of the current node's left node
left_node = node.left
if left_node is None:
return
else:
# left node's right node become the left node of current node
node.left = left_node.right
if left_node.right is not None:
left_node.right.parent = node
left_node.parent = node.parent
# check the parent case
if node.parent is None:
self.root = left_node
elif node is node.parent.left:
node.parent.left = left_node
else:
node.parent.right = left_node
left_node.right = node
node.parent = left_node
def insert(self, node):
# the inserted node's color is default is red
root = self.root
insert_node_parent = None
# find the position of inserted node
while root is not None:
insert_node_parent = root
if insert_node_parent.val < node.val:
root = root.right
else:
root = root.left
# set the n ode's parent node
node.parent = insert_node_parent
if insert_node_parent is None:
# case 1 inserted tree is null
self.root = node
elif insert_node_parent.val > node.val:
# case 2 not null and find left or right
insert_node_parent.left = node
else:
insert_node_parent.right = node
node.left = None
node.right = None
node.color = 1
# fix the tree to
self.fix_insert(node)
def fix_insert(self, node):
# case 1 the parent is null, then set the inserted node as root and color = 0
if node.parent is None:
node.color = 0
self.root = node
return
# case 2 the parent color is black, do nothing
# case 3 the parent color is red
while node.parent and node.parent.color is 1:
if node.parent is node.parent.parent.left:
uncle_node = node.parent.parent.right
if uncle_node and uncle_node.color is 1:
# case 3.1 the uncle node is red
# then set parent and uncle color is black and grandparent is red
# then node => node.parent
node.parent.color = 0
node.parent.parent.right.color = 0
node.parent.parent.color = 1
node = node.parent.parent
continue
elif node is node.parent.right:
# case 3.2 the uncle node is black or null, and the node is right of parent
# then set his parent node is current node
# left rotate the node and continue the next
node = node.parent
self.left_rotate(node)
# case 3.3 the uncle node is black and parent node is left
# then parent node set black and grandparent set red
node.parent.color = 0
node.parent.parent.color = 1
self.right_rotate(node.parent.parent)
else:
uncle_node = node.parent.parent.left
if uncle_node and uncle_node.color is 1:
# case 3.1 the uncle node is red
# then set parent and uncle color is black and grandparent is red
# then node => node.parent
node.parent.color = 0
node.parent.parent.left.color = 0
node.parent.parent.color = 1
node = node.parent.parent
continue
elif node is node.parent.left:
# case 3.2 the uncle node is black or null, and the node is right of parent
# then set his parent node is current node
# left rotate the node and continue the next
node = node.parent
self.right_rotate(node)
# case 3.3 the uncle node is black and parent node is left
# then parent node set black and grandparent set red
node.parent.color = 0
node.parent.parent.color = 1
self.left_rotate(node.parent.parent)
self.root.color = 0
def transplant(self, node_u, node_v):
"""
replace u with v
:param node_u: replaced node
:param node_v:
:return: None
"""
if node_u.parent is None:
self.root = node_v
elif node_u is node_u.parent.left:
node_u.parent.left = node_v
elif node_u is node_u.parent.right:
node_u.parent.right = node_v
# check is node_v is None
if node_v:
node_v.parent = node_u.parent
def maximum(self, node):
"""
find the max node when node regard as a root node
:param node:
:return: max node
"""
temp_node = node
while temp_node.right is not None:
temp_node = temp_node.right
return temp_node
def minimum(self, node):
"""
find the minimum node when node regard as a root node
:param node:
:return: minimum node
"""
temp_node = node
while temp_node.left:
temp_node = temp_node.left
return temp_node
def delete(self, node):
# find the node position
node_color = node.color
if node.left is None:
temp_node = node.right
self.transplant(node, node.right)
elif node.right is None:
temp_node = node.left
self.transplant(node, node.left)
else:
# both child exits ,and find minimum child of right child
node_min = self.minimum(node.right)
node_color = node_min.color
temp_node = node_min.right
##
if node_min.parent != node:
self.transplant(node_min, node_min.right)
node_min.right = node.right
node_min.right.parent = node_min
self.transplant(node, node_min)
node_min.left = node.left
node_min.left.parent = node_min
node_min.color = node.color
# when node is black ,then need to fix it with 4 cases
if node_color == 0:
self.delete_fixup(temp_node)
def delete_fixup(self, node):
# 4 cases
while node != self.root and node.color == 0:
# node is not root and color is black
if node == node.parent.left:
# node is left node
node_brother = node.parent.right
# case 1: node's red, can not get black node
# set brother is black and parent is red
if node_brother.color == 1:
node_brother.color = 0
node.parent.color = 1
self.left_rotate(node.parent)
node_brother = node.parent.right
# case 2: brother node is black, and its children node is both black
if (node_brother.left is None or node_brother.left.color == 0) and (
node_brother.right is None or node_brother.right.color == 0):
node_brother.color = 1
node = node.parent
else:
# case 3: brother node is black , and its left child node is red and right is black
if node_brother.right is None or node_brother.right.color == 0:
node_brother.color = 1
node_brother.left.color = 0
self.right_rotate(node_brother)
node_brother = node.parent.right
# case 4: brother node is black, and right is red, and left is any color
node_brother.color = node.parent.color
node.parent.color = 0
node_brother.right.color = 0
self.left_rotate(node.parent)
node = self.root
break
else:
node_brother = node.parent.left
if node_brother.color == 1:
node_brother.color = 0
node.parent.color = 1
self.left_rotate(node.parent)
node_brother = node.parent.right
if (node_brother.left is None or node_brother.left.color == 0) and (
node_brother.right is None or node_brother.right.color == 0):
node_brother.color = 1
node = node.parent
else:
if node_brother.left is None or node_brother.left.color == 0:
node_brother.color = 1
node_brother.right.color = 0
self.left_rotate(node_brother)
node_brother = node.parent.left
node_brother.color = node.parent.color
node.parent.color = 0
node_brother.left.color = 0
self.right_rotate(node.parent)
node = self.root
break
node.color = 0
def inorder(self):
res = []
if not self.root:
return res
stack = []
root = self.root
while root or stack:
while root:
stack.append(root)
root = root.left
root = stack.pop()
res.append({'val': root.val, 'color': root.color})
root = root.right
return res
if __name__ == "__main__":
rb = RBTree()
children = [11, 2, 14, 1, 7, 15, 5, 8, 4]
for child in children:
node = RBNode(child, 1)
print(child)
rb.insert(node)
print(rb.inorder())