-
Notifications
You must be signed in to change notification settings - Fork 47
/
red_black_tree.rb
506 lines (441 loc) · 16.2 KB
/
red_black_tree.rb
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
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
:BLACK
:RED
:NIL_C
:LEFT
:RIGHT
# Contain combinations of directions
DIRECTIONS = {
[:LEFT, :RIGHT] => 'LR',
[:RIGHT, :LEFT] => 'RL',
[:RIGHT, :RIGHT] => 'RR',
[:LEFT, :LEFT] => 'LL'
}
class Node
attr_accessor :value, :color, :parent, :left, :right
def initialize(value, color, parent, left, right)
@value = value
@color = color
@parent = parent
@left = left
@right = right
end
def ==(other)
if self.color == :NIL_C && other.color==self.color
return true
end
same_parents = false
if self.parent.nil? || other.parent.nil?
same_parents = self.parent.nil? && other.parent.nil?
else
same_parents = self.parent.value == other.parent.value
end
self.value == other.value && self.color == other.color && same_parents
end
def has_children?
self.get_children_count != 0
end
def get_children_count
if self.color == :NIL_C
return 0
end
if self.left.color != :NIL_C && self.right.color != :NIL_C
return 2
elsif self.left.color != :NIL_C || self.right.color != :NIL_C
return 1
else
return 0
end
end
def to_s
"#{@color} #{@value}"
end
def inspect
to_s
end
end
class RedBlackTree
@@nil_leaf = Node.new(value=NIL, color=:NIL_C, parent=NIL, left=NIL, right=NIL)
attr_accessor :count, :root
def initialize
@count = 0
@root = NIL
end
def add(value)
if @root.nil?
@root = Node.new(value=value, color=:BLACK, parent=NIL, left=@@nil_leaf, right=@@nil_leaf)
@count += 1
return
end
parent, direction = find_parent value
if direction.nil?
return # value is in the tree
end
new_node = Node.new(value=value, color=:RED, parent=parent, left=@@nil_leaf, right=@@nil_leaf)
if direction == :LEFT
parent.left = new_node
else
parent.right = new_node
end
try_rebalance new_node
@count += 1
end
def try_rebalance(new_node)
parent = new_node.parent
value = new_node.value
if parent.nil? || # what the fuck?
parent.parent.nil? || # parent is root
parent.color != :RED # no red-red, no problem :)
return
end
grandfather = parent.parent
direction = if parent.value > new_node.value then :LEFT else :RIGHT end
parent_direction = if grandfather.value > parent.value then :LEFT else :RIGHT end
uncle = if parent_direction == :LEFT then grandfather.right else grandfather.left end
general_direction = DIRECTIONS[[direction, parent_direction]]
if uncle == @@nil_leaf || uncle.color == :BLACK
if general_direction == 'LL'
# LL => Right rotation
right_rotation(new_node, parent, grandfather, to_recolor=true)
elsif general_direction == 'RR'
# RR => Left rotation
left_rotation(new_node, parent, grandfather, to_recolor=true)
elsif general_direction == 'LR'
# LR => Right rotation, left rotation
right_rotation(NIL, new_node, parent, to_recolor=false)
# due to the right rotation, parent and new_node positions have switched
left_rotation(parent, new_node, grandfather, to_recolor=true)
elsif general_direction == 'RL'
# RL => Left rotation, right rotation
left_rotation(NIL, new_node, parent, to_recolor=false)
# due to the left rotation, parent and new_node positions have switches
right_rotation(parent, new_node, grandfather, to_recolor=true)
end
else
# uncle is red, simply recolor
recolor(grandfather)
end
end
def right_rotation(node, parent, grandfather, to_recolor=false)
grand_grandfather = grandfather.parent
# grandfather will become the right child of parent
update_parent(node=parent, parent_old_child=grandfather, new_parent=grand_grandfather)
old_right = parent.right
parent.right = grandfather
grandfather.parent = parent
grandfather.left = old_right
old_right.parent = grandfather
if to_recolor # recolor the nodes after a move to preserve invariants
parent.color = :BLACK
node.color = :RED
grandfather.color = :RED
end
end
def left_rotation(node, parent, grandfather, to_recolor=false)
grand_grandfather = grandfather.parent
# grandfather will become the left child of parent
update_parent(node=parent, parent_old_child=grandfather, new_parent=grand_grandfather)
old_left = parent.left
parent.left = grandfather
grandfather.parent = parent
grandfather.right = old_left
old_left.parent = grandfather
if to_recolor
parent.color = :BLACK
grandfather.color = :RED
node.color = :RED
end
end
# recolors the grandfather red, coloring his children black
def recolor(grandfather)
grandfather.left.color = :BLACK
grandfather.right.color = :BLACK
if @root != grandfather
grandfather.color = :RED
end
try_rebalance grandfather
end
# our node 'switches' place with the old child, assigning a new parent to the node
# if the new_parent is NIL, this means that our node becomes the root of the tree
def update_parent(node, parent_old_child, new_parent)
node.parent = new_parent
if not new_parent.nil?
# determine the old child's position to put the node there
if new_parent.value > parent_old_child.value
new_parent.left = node
else
new_parent.right = node
end
else
@root = node
end
end
# finds a place for the value in the binary tree, returning the node and the direction it should go in
def find_parent(value)
find = lambda do |node|
if node.value == value
return
elsif node.value > value
# go left
if node.left.color == :NIL_C
# no more to go
return node, :LEFT
end
return find.call node.left
else
# go right
if node.right.color == :NIL_C
# no more to go
return node, :RIGHT
end
return find.call node.right
end
end
find.call @root
end
# removes a node from the tree
def remove(value)
# try to get a node with 0 or 1 children (by finding a successor if needed)
node_to_remove = find_node value
if node_to_remove.nil?
return # value not in tree
end
if node_to_remove.get_children_count == 2 # does not have 0 or 1 children, find successor
# find the in-order successor and replace its value, then remove the successor
successor = find_successor node_to_remove
node_to_remove.value = successor.value
node_to_remove = successor
end
remove_internal node_to_remove
@count -= 1
end
# receives a node with 0 or 1 children and removes it according to its color/children
def remove_internal(node_to_remove)
left_child = node_to_remove.left
right_child = node_to_remove.right
not_nil_child = if left_child != @@nil_leaf then left_child else right_child end
if node_to_remove == self.root
if not_nil_child != @@nil_leaf
@root = not_nil_child
@root.parent = NIL
@root.color = :BLACK
else # both children are nil and this is the root
@root = NIL
return
end
if node_to_remove.color == :RED
# red node with no children, simple remove
if node_to_remove.has_children?
raise Exception('Unexpected behavior, a successor red node without 2 children should not have any children!')
end
remove_leaf node
else # node is black
if not_nil_child.color == :RED
# last easy chance, swap the values with the red child and simply remove it
node_to_remove.value = not_nil_child.value
node_to_remove.left = not_nil_child.left
node_to_remove.right = not_nil_child.right
else # black child
# 6 different cases apply here, good luck
remove_black_node node_to_remove
end
end
end
end
# loop through each of the 6 cases until we reach a terminating case
# what we're left is a leaf node ready to be deleted
def remove_black_node(node_to_remove)
# code here
case_1 node_to_remove
remove_leaf node_to_remove
end
# simply removes a leaf node, making its parent point to a NIL_LEAF
def remove_leaf(node)
if node.parent.value > node.value
node.parent.left = @@nil_leaf
else
node.parent.right = @@nil_leaf
end
end
def find_successor(node_to_remove)
successor = node_to_remove.right
successor = successor.left while successor.left != @@nil_leaf
successor
end
def find_node(value)
find = Proc.new do |node|
if node.nil? || node == @@nil_leaf
return NIL
elsif node.value == value
return node
elsif node.value > value
return find.call node.left
else
return find.call node.right
end
end
find.call @root
end
# Case 1 is when there's a double black node on the root
# Because we're at the root, we can simply remove it
# and reduce the black height of the whole tree.
#
# __|10B|__ __10B__
# / \ ==> / \
#9B 20B 9B 20B
def case_1(node_to_remove)
if @root == node_to_remove
node_to_remove.color = :BLACK
return
end
case_2 node_to_remove
end
# case 2 applies when
# the parent is BLACK
# the sibling is RED
# the sibling's children are BLACK or NIL
#It takes the sibling and rotates it
#
# 40B 60B
# / \ --CASE 2 ROTATE--> / \
# |20B| 60R LEFT ROTATE 40R 80B
# DBL BLACK IS 20----^ / \ SIBLING 60R / \
# 50B 80B |20B| 50B
# (if the sibling's direction was left of its parent, we would RIGHT ROTATE it)
# Now the original node's parent is RED
# and we can apply case 4 or case 6
def case_2(node_to_remove)
parent = node_to_remove.parent
sibling, direction = get_sibling node_to_remove
if sibling.color == :RED && parent.color == :BLACK && sibling.left.color != :RED && sibling.right.color != :RED
if direction == :RIGHT
left_rotation(node=NIL, parent=sibling, grandfather=parent)
else
right_rotation(node=NIL, parent=sibling, grandfather=parent)
end
parent.color = :RED
sibling.color = :BLACK
return case_1 node_to_remove
end
case_3 node_to_remove
end
#Case 3 deletion is when:
# the parent is BLACK
# the sibling is BLACK
# the sibling's children are BLACK (or nil)
# Then, we make the sibling red and
# pass the double black node upwards
#
# Parent is black
# ___50B___ Sibling is black ___50B___
# / \ Sibling's children are black / \
# 30B 80B CASE 3 30B |80B| Continue with other cases
# / \ / \ ==> / \ / \
# 20B 35R 70B |90B|<---REMOVE 20B 35R 70R X
# / \ / \
# 34B 37B 34B 37B
def case_3(node_to_remove)
parent = node_to_remove.parent
sibling, _ = get_sibling node_to_remove
if sibling.color == :BLACK && parent.color == :BLACK && sibling.left.color != :RED && sibling.right.color != :RED
# color the sibling red and forward the black node upwards, calling the cases for the parent
sibling.color = :RED
return case_1 parent
end
case_4 node_to_remove
end
# TERMINATING CASE
# If the parent is red and the sibling is black with no red children,
# simply swap their colors
# DB-Double Black
# __10R__ __10B__ The black height of the left subtree has been incremented
# / \ / \ And the one below stays the same
# DB 15B ===> X 15R No consequences, we're done!
# / \ / \
# 12B 17B 12B 17B
def case_4(node_to_remove)
parent = node_to_remove.parent
sibling, _ = get_sibling node_to_remove
if parent.color == :RED && sibling.color == :BLACK && sibling.left.color != :RED && sibling.right.color != :RED
parent.color = :BLACK
sibling.color = :RED
return
end
case_5 node_to_remove
end
# Case 5 is a rotation that changes the circumstances so that we can do a case 6
# If the closer node is red and the outer BLACK or NIL, we do a left/right rotation, depending on the orientation
# This showcases when the CLOSER NODE's direction is RIGHT
#
# ___50B___ __50B__
# / \ / \
# 30B |80B| <-- Double black 35B |80B| Case 6 is now
# / \ / \ Closer node is red (35R) / \ / applicable here,
# 20B 35R 70R X Outer is black (20B) 30R 37B 70R so we redirect the node
# / \ So we do a LEFT ROTATION / \ to it :)
# 34B 37B on 35R (closer node) 20B 34B
def case_5(node_to_remove)
sibling, direction = get_sibling node_to_remove
closer_node = if direction == :LEFT then sibling.right else sibling.left end
outer_node = if direction == :LEFT then sibling.left else sibling.right end
if closer_node.color == :RED && outer_node.color != :RED && sibling.color == :BLACK
if direction == :LEFT
left_rotation(node=NIL, parent=closer_node, grandfather=sibling)
else
right_rotation(node=NIL, parent=closer_node, grandfather=sibling)
end
closer_node.color = :BLACK
sibling.color = :RED
end
case_6 node_to_remove
end
# TERMINATING
# Case 6 requires
# SIBLING to be BLACK
# OUTER NODE to be RED
# Then, does a right or left rotation on the sibling
# This will showcase when the SIBLING's direction is LEFT
#
# Double Black
# __50B__ | __35B__
# / \ | / \
# SIBLING--> 35B |80B| <- 30R 50R
# / \ / / \ / \
# 30R 37B 70R Outer node is RED 20B 34B 37B 80B
# / \ Closer node doesn't /
# 20B 34B matter 70R
# Parent doesn't
# matter
# So we do a right rotation on 35B!
def case_6(node_to_remove)
sibling, direction = get_sibling node_to_remove
outer_node = if direction == :LEFT then sibling.left else sibling.right end
case_6_rotation = Proc.new do |dir|
parent_color = sibling.parent.color
if dir == :LEFT
right_rotation(node=NIL, parent=sibling, grandfather=sibling.parent)
else
left_rotation(node=NIL, parent=sibling, grandfather=sibling.parent)
end
# our new parent is the sibling
sibling.color = parent_color
sibling.right.color = :BLACK
sibling.left.color = :BLACK
end
if sibling.color == :BLACK && outer_node.color == :RED
return case_6_rotation(direction) # terminating
end
raise Exception('Should not have reached here')
end
def get_sibling(node)
# code here
parent = node.parent
parent_sibling, dir = if parent.left == node then [parent.right, :RIGHT] else [parent.left, :LEFT] end
return parent_sibling, dir
end
public :add
private :find_parent, :update_parent, :recolor, :right_rotation, :left_rotation
# end
def self.get_nil_leaf
return @@nil_leaf
end
end