fix(rbt): use explicit colours instead of constants

Signed-off-by: Matej Focko <mfocko@redhat.com>
This commit is contained in:
Matej Focko 2022-05-20 15:47:03 +02:00
parent d39e11713c
commit 2b02eac10e
Signed by: mfocko
GPG key ID: 7C47D46246790496

14
rbt.py
View file

@ -61,7 +61,7 @@ class RBTree(RankedTree[T]):
# Case 1
# ======
# zs uncle y is red
pp.rank += 1
pp.rank += Colour.Black
z = pp
elif z == right_child:
# Case 2
@ -105,33 +105,33 @@ class RBTree(RankedTree[T]):
rotate_left: RotateFunction,
rotate_right: RotateFunction,
) -> Tuple[Optional[Node[T]], Optional[Node[T]]]:
if Node.difference(w) == 0:
if Node.difference(w) == Colour.Red:
# Case 1
# ======
# xs sibling w is red
rotate_left(tree=self, x=parent)
w = right(parent)
if Node.differences(w) == (1, 1):
if Node.differences(w) == (Colour.Black, Colour.Black):
# Case 2
# ======
# xs sibling w is black, and both of ws children are black
parent.rank -= 1
parent.rank -= Colour.Black
x = parent
else:
# Case 3
# ======
# xs sibling w is black,
# ws left child is red, and ws right child is black
if Node.difference(right(w), w) == 1:
if Node.difference(right(w), w) == Colour.Black:
rotate_right(tree=self, x=w)
w = right(parent)
# Case 4
# ======
# xs sibling w is black, and ws right child is red
parent.rank -= 1
w.rank += 1
parent.rank -= Colour.Black
w.rank += Colour.Black
rotate_left(tree=self, x=parent)
x = self.root