Skip to content

Commit 546e402

Browse files
AlexeyBasingerAlexeyBasinger
committed
Refactor red-black tree removal to reduce cyclomatic complexity
1 parent a051ab5 commit 546e402

File tree

1 file changed

+254
-113
lines changed

1 file changed

+254
-113
lines changed

‎data_structures/binary_tree/red_black_tree.py‎

Lines changed: 254 additions & 113 deletions
Original file line numberDiff line numberDiff line change
@@ -112,10 +112,8 @@ def insert(self, label: int) -> RedBlackTree:
112112
def _insert_repair(self) -> None:
113113
"""Repair the coloring from inserting into a tree."""
114114
if self.parent is None:
115-
# This node is the root, so it just needs to be black
116115
self.color = 0
117116
elif color(self.parent) == 0:
118-
# If the parent is black, then it just needs to be red
119117
self.color = 1
120118
else:
121119
uncle = self.parent.sibling
@@ -147,133 +145,276 @@ def _insert_repair(self) -> None:
147145
self.grandparent.color = 1
148146
self.grandparent._insert_repair()
149147

150-
def remove(self, label: int) -> RedBlackTree:
151-
"""Remove label from this tree."""
152-
if self.label == label:
153-
if self.left and self.right:
154-
# It's easier to balance a node with at most one child,
155-
# so we replace this node with the greatest one less than
156-
# it and remove that.
157-
value = self.left.get_max()
158-
if value is not None:
159-
self.label = value
160-
self.left.remove(value)
161-
else:
162-
# This node has at most one non-None child, so we don't
163-
# need to replace
164-
child = self.left or self.right
165-
if self.color == 1:
166-
# This node is red, and its child is black
167-
# The only way this happens to a node with one child
168-
# is if both children are None leaves.
169-
# We can just remove this node and call it a day.
170-
if self.parent:
171-
if self.is_left():
172-
self.parent.left = None
173-
else:
174-
self.parent.right = None
175-
# The node is black
176-
elif child is None:
177-
# This node and its child are black
178-
if self.parent is None:
179-
# The tree is now empty
180-
return RedBlackTree(None)
181-
else:
182-
self._remove_repair()
183-
if self.is_left():
184-
self.parent.left = None
185-
else:
186-
self.parent.right = None
187-
self.parent = None
188-
else:
189-
# This node is black and its child is red
190-
# Move the child node here and make it black
191-
self.label = child.label
192-
self.left = child.left
193-
self.right = child.right
194-
if self.left:
195-
self.left.parent = self
196-
if self.right:
197-
self.right.parent = self
198-
elif self.label is not None and self.label > label:
199-
if self.left:
200-
self.left.remove(label)
201-
elif self.right:
202-
self.right.remove(label)
148+
def remove(self, label: int) -> "RedBlackTree":
149+
"""Remove label from this tree"""
150+
if self.label is None:
151+
return self
152+
153+
target = self.search(label)
154+
if target is None:
155+
return self.parent or self
156+
157+
target._remove_node()
203158
return self.parent or self
204159

160+
def _remove_node(self) -> None:
161+
"""
162+
Physically remove the current node from the tree,
163+
preserving all the invariants of the red-black tree
164+
"""
165+
if self.left and self.right:
166+
self._remove_with_two_children()
167+
else:
168+
self._remove_with_zero_or_one_child()
169+
170+
def _remove_with_two_children(self) -> None:
171+
"""
172+
Handling the case when a node has two non-empty children:
173+
Find the maximum in the left subtree, copy the value,
174+
Delete the maximum in the left subtree
175+
"""
176+
left = self.left
177+
if left is None:
178+
# Logically this should not happen if the caller knows that
179+
# the node has two children, but this guard keeps type
180+
# checkers happy and makes the method safer.
181+
return
182+
183+
value = left.get_max()
184+
if value is None:
185+
# No value in the left subtree – nothing to do.
186+
return
187+
188+
# Copy the predecessor value into the current node and
189+
# delete the predecessor from the left subtree.
190+
self.label = value
191+
left.remove(value)
192+
193+
def _remove_with_zero_or_one_child(self) -> None:
194+
"""
195+
Handling the case when a node has 0 or 1 child
196+
"""
197+
child = self.left or self.right
198+
199+
if self.color == 1:
200+
self._remove_red_leaf()
201+
return
202+
203+
if child is None:
204+
self._remove_black_leaf()
205+
return
206+
207+
self._remove_black_node_with_red_child(child)
208+
209+
def _remove_red_leaf(self) -> None:
210+
"""
211+
delete red leaf
212+
"""
213+
if self.parent is None:
214+
self.label = None
215+
return
216+
217+
if self.is_left():
218+
self.parent.left = None
219+
else:
220+
self.parent.right = None
221+
self.parent = None
222+
223+
def _remove_black_leaf(self) -> None:
224+
"""
225+
delete black leaf
226+
"""
227+
if self.parent is None:
228+
self.label = None
229+
return
230+
231+
self._remove_repair()
232+
233+
if self.is_left():
234+
self.parent.left = None
235+
else:
236+
self.parent.right = None
237+
self.parent = None
238+
239+
def _remove_black_node_with_red_child(self, child: "RedBlackTree") -> None:
240+
"""
241+
Black knot with a single red child:
242+
Move the child to the top and paint it black.
243+
"""
244+
self.label = child.label
245+
self.left = child.left
246+
self.right = child.right
247+
if self.left:
248+
self.left.parent = self
249+
if self.right:
250+
self.right.parent = self
251+
self.color = 0
252+
205253
def _remove_repair(self) -> None:
206-
"""Repair the coloring of the tree that may have been messed up."""
254+
"""Repair the coloring of the tree that may have been messed up
255+
after deleting a black node.
256+
"""
207257
if (
208258
self.parent is None
209259
or self.sibling is None
210-
or self.parent.sibling is None
211-
or self.grandparent is None
260+
or self.parent.grandparent is None
212261
):
213262
return
214-
if color(self.sibling) == 1:
215-
self.sibling.color = 0
216-
self.parent.color = 1
217-
if self.is_left():
218-
self.parent.rotate_left()
219-
else:
220-
self.parent.rotate_right()
221-
if (
222-
color(self.parent) == 0
223-
and color(self.sibling) == 0
224-
and color(self.sibling.left) == 0
225-
and color(self.sibling.right) == 0
226-
):
227-
self.sibling.color = 1
228-
self.parent._remove_repair()
263+
264+
self._repair_red_sibling()
265+
266+
if self._repair_black_parent_black_sibling_black_children():
229267
return
230-
if (
231-
color(self.parent) == 1
232-
and color(self.sibling) == 0
233-
and color(self.sibling.left) == 0
234-
and color(self.sibling.right) == 0
235-
):
236-
self.sibling.color = 1
237-
self.parent.color = 0
268+
269+
if self._repair_red_parent_black_sibling_black_children():
238270
return
271+
272+
self._repair_inner_nephew()
273+
274+
self._repair_outer_nephew()
275+
276+
def _repair_red_sibling(self) -> None:
277+
"""Case 1: sibling is red.
278+
279+
We rotate around the parent so that the sibling becomes black,
280+
and then we continue with a configuration where the sibling is
281+
black and the parent is red.
282+
"""
283+
sibling = self.sibling
284+
parent = self.parent
285+
286+
if sibling is None or parent is None:
287+
return
288+
289+
if color(sibling) != 1:
290+
return
291+
292+
sibling.color = 0
293+
parent.color = 1
294+
295+
if self.is_left():
296+
parent.rotate_left()
297+
else:
298+
parent.rotate_right()
299+
300+
def _repair_black_parent_black_sibling_black_children(self) -> bool:
301+
"""Case 2:
302+
parent black, sibling black, sibling.left & sibling.right black.
303+
304+
In this case we recolor the sibling red and propagate the
305+
"double black" upwards to the parent.
306+
"""
307+
parent = self.parent
308+
sibling = self.sibling
309+
310+
if parent is None or sibling is None:
311+
return False
312+
313+
if color(parent) != 0:
314+
return False
315+
if color(sibling) != 0:
316+
return False
317+
if color(sibling.left) != 0:
318+
return False
319+
if color(sibling.right) != 0:
320+
return False
321+
322+
sibling.color = 1
323+
parent._remove_repair()
324+
return True
325+
326+
def _repair_red_parent_black_sibling_black_children(self) -> bool:
327+
"""Case 3:
328+
parent red, sibling black, sibling.left & sibling.right black.
329+
330+
We just swap the colors of the parent and sibling and finish.
331+
"""
332+
parent = self.parent
333+
sibling = self.sibling
334+
335+
if parent is None or sibling is None:
336+
return False
337+
338+
if color(parent) != 1:
339+
return False
340+
if color(sibling) != 0:
341+
return False
342+
if color(sibling.left) != 0:
343+
return False
344+
if color(sibling.right) != 0:
345+
return False
346+
347+
sibling.color = 1
348+
parent.color = 0
349+
return True
350+
351+
def _repair_inner_nephew(self) -> None:
352+
"""Case 4: inner nephew is red.
353+
354+
We rotate around the sibling to turn this into the outer-nephew
355+
case (case 5), which can then be fixed by a rotation around
356+
the parent.
357+
"""
358+
sibling = self.sibling
359+
if sibling is None:
360+
return
361+
362+
# Left child, red left (inner) nephew
239363
if (
240364
self.is_left()
241-
and color(self.sibling) == 0
242-
and color(self.sibling.right) == 0
243-
and color(self.sibling.left) == 1
244-
):
245-
self.sibling.rotate_right()
246-
self.sibling.color = 0
247-
if self.sibling.right:
248-
self.sibling.right.color = 1
249-
if (
250-
self.is_right()
251-
and color(self.sibling) == 0
252-
and color(self.sibling.right) == 1
253-
and color(self.sibling.left) == 0
254-
):
255-
self.sibling.rotate_left()
256-
self.sibling.color = 0
257-
if self.sibling.left:
258-
self.sibling.left.color = 1
259-
if (
260-
self.is_left()
261-
and color(self.sibling) == 0
262-
and color(self.sibling.right) == 1
365+
and color(sibling) == 0
366+
and color(sibling.right) == 0
367+
and color(sibling.left) == 1
263368
):
264-
self.parent.rotate_left()
265-
self.grandparent.color = self.parent.color
266-
self.parent.color = 0
267-
self.parent.sibling.color = 0
369+
sibling.rotate_right()
370+
sibling.color = 0
371+
if sibling.right:
372+
sibling.right.color = 1
373+
374+
# Right child, red right (inner) nephew
268375
if (
269376
self.is_right()
270-
and color(self.sibling) == 0
271-
and color(self.sibling.left) == 1
377+
and color(sibling) == 0
378+
and color(sibling.right) == 1
379+
and color(sibling.left) == 0
272380
):
273-
self.parent.rotate_right()
274-
self.grandparent.color = self.parent.color
275-
self.parent.color = 0
276-
self.parent.sibling.color = 0
381+
sibling.rotate_left()
382+
sibling.color = 0
383+
if sibling.left:
384+
sibling.left.color = 1
385+
386+
def _repair_outer_nephew(self) -> None:
387+
"""Case 5: outer nephew is red.
388+
389+
This is the final case: a rotation around the parent and
390+
recoloring of parent / sibling / grandparent fixes the violation.
391+
"""
392+
sibling = self.sibling
393+
parent = self.parent
394+
grandparent = self.grandparent
395+
396+
if sibling is None or parent is None or grandparent is None:
397+
return
398+
399+
# Left child, red right (outer) nephew
400+
if self.is_left() and color(sibling) == 0 and color(sibling.right) == 1:
401+
parent.rotate_left()
402+
grandparent.color = parent.color
403+
parent.color = 0
404+
405+
parent_sibling = parent.sibling
406+
if parent_sibling is not None:
407+
parent_sibling.color = 0
408+
409+
# Right child, red left (outer) nephew
410+
if self.is_right() and color(sibling) == 0 and color(sibling.left) == 1:
411+
parent.rotate_right()
412+
grandparent.color = parent.color
413+
parent.color = 0
414+
415+
parent_sibling = parent.sibling
416+
if parent_sibling is not None:
417+
parent_sibling.color = 0
277418

278419
def check_color_properties(self) -> bool:
279420
"""Check the coloring of the tree, and return True iff the tree

0 commit comments

Comments
 (0)