2022-05-22 21:00:18 +02:00
|
|
|
|
from node import Node, Comparable
|
|
|
|
|
from ranked_tree import RankedTree, RotateFunction
|
2022-05-20 00:01:04 +02:00
|
|
|
|
|
|
|
|
|
import enum
|
|
|
|
|
import logging
|
|
|
|
|
from typing import Callable, Tuple, TypeVar, Optional
|
|
|
|
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
T = TypeVar("T", bound=Comparable)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class Colour(enum.IntEnum):
|
2022-05-22 00:38:48 +02:00
|
|
|
|
"""
|
|
|
|
|
Represents colour of the edge or node.
|
|
|
|
|
"""
|
2022-05-22 21:00:18 +02:00
|
|
|
|
|
2022-05-20 00:01:04 +02:00
|
|
|
|
Red = 0
|
|
|
|
|
Black = 1
|
|
|
|
|
|
|
|
|
|
|
2022-05-22 21:07:40 +02:00
|
|
|
|
def has_double_black(x: Optional[Node[T]]) -> bool:
|
2022-05-22 00:38:48 +02:00
|
|
|
|
"""
|
|
|
|
|
Checks for double black child of x.
|
|
|
|
|
|
|
|
|
|
Args:
|
|
|
|
|
x: Node to be checked.
|
|
|
|
|
|
|
|
|
|
Returns:
|
|
|
|
|
`true`, if `x` has a double black node, `false` otherwise.
|
|
|
|
|
"""
|
|
|
|
|
return x is not None and 2 in Node.differences(x)
|
2022-05-20 00:01:04 +02:00
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class RBTree(RankedTree[T]):
|
|
|
|
|
def is_correct_node(
|
|
|
|
|
self, node: Optional[Node[T]], recursive: bool = True
|
|
|
|
|
) -> bool:
|
|
|
|
|
if not node:
|
|
|
|
|
return True
|
|
|
|
|
|
|
|
|
|
left, right = Node.differences(node)
|
|
|
|
|
if left not in (Colour.Red, Colour.Black):
|
|
|
|
|
# left subtree has invalid difference
|
|
|
|
|
return False
|
|
|
|
|
elif right not in (Colour.Red, Colour.Black):
|
|
|
|
|
# right subtree has invalid difference
|
|
|
|
|
return False
|
|
|
|
|
|
|
|
|
|
if Node.difference(node) == Colour.Red and (
|
|
|
|
|
left == Colour.Red or right == Colour.Red
|
|
|
|
|
):
|
|
|
|
|
# two consecutive red nodes
|
|
|
|
|
return False
|
|
|
|
|
|
|
|
|
|
return not recursive or (
|
|
|
|
|
self.is_correct_node(node.left)
|
|
|
|
|
and self.is_correct_node(node.right)
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
# region InsertRebalance
|
|
|
|
|
|
|
|
|
|
def _insert_rebalance_step(
|
|
|
|
|
self,
|
2022-05-20 11:47:38 +02:00
|
|
|
|
z: Node[T],
|
|
|
|
|
y: Optional[Node[T]],
|
|
|
|
|
right_child: Node[T],
|
2022-05-22 21:07:40 +02:00
|
|
|
|
rotate_left: RotateFunction[T],
|
|
|
|
|
rotate_right: RotateFunction[T],
|
2022-05-20 00:01:04 +02:00
|
|
|
|
) -> Node[T]:
|
2022-05-20 11:47:38 +02:00
|
|
|
|
p = z.parent
|
2022-05-20 00:01:04 +02:00
|
|
|
|
pp = p.parent
|
|
|
|
|
|
2022-05-20 11:47:38 +02:00
|
|
|
|
if y is not None and Node.difference(y) == Colour.Red:
|
|
|
|
|
# Case 1
|
|
|
|
|
# ======
|
|
|
|
|
# z’s uncle y is red
|
2022-05-20 15:47:03 +02:00
|
|
|
|
pp.rank += Colour.Black
|
2022-05-20 11:47:38 +02:00
|
|
|
|
z = pp
|
|
|
|
|
elif z == right_child:
|
|
|
|
|
# Case 2
|
|
|
|
|
# ======
|
|
|
|
|
# z’s uncle y is black and z is a right child
|
|
|
|
|
z = p
|
2022-05-22 21:00:18 +02:00
|
|
|
|
rotate_left(p)
|
2022-05-20 00:01:04 +02:00
|
|
|
|
else:
|
2022-05-20 11:47:38 +02:00
|
|
|
|
# Case 3
|
|
|
|
|
# ======
|
|
|
|
|
# z’s uncle y is black and z is a left child
|
2022-05-22 21:00:18 +02:00
|
|
|
|
rotate_right(pp)
|
2022-05-20 00:01:04 +02:00
|
|
|
|
|
2022-05-20 11:47:38 +02:00
|
|
|
|
return z
|
2022-05-20 00:01:04 +02:00
|
|
|
|
|
2022-05-20 11:47:38 +02:00
|
|
|
|
def _insert_rebalance(self, z: Node[T]) -> None:
|
|
|
|
|
while z.parent is not None and Node.difference(z.parent) == Colour.Red:
|
|
|
|
|
p = z.parent
|
|
|
|
|
pp = p.parent
|
2022-05-20 00:01:04 +02:00
|
|
|
|
|
|
|
|
|
assert pp
|
|
|
|
|
if p == pp.left:
|
2022-05-20 11:47:38 +02:00
|
|
|
|
z = self._insert_rebalance_step(
|
2022-05-22 21:00:18 +02:00
|
|
|
|
z, pp.right, p.right, self.rotate_left, self.rotate_right
|
2022-05-20 00:01:04 +02:00
|
|
|
|
)
|
|
|
|
|
else:
|
2022-05-20 11:47:38 +02:00
|
|
|
|
z = self._insert_rebalance_step(
|
2022-05-22 21:00:18 +02:00
|
|
|
|
z, pp.left, p.left, self.rotate_right, self.rotate_left
|
2022-05-20 00:01:04 +02:00
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
# endregion InsertRebalance
|
|
|
|
|
|
|
|
|
|
# region DeleteRebalance
|
|
|
|
|
|
|
|
|
|
def _delete_rebalance_step(
|
|
|
|
|
self,
|
|
|
|
|
x: Node[T],
|
|
|
|
|
w: Node[T],
|
|
|
|
|
parent: Node[T],
|
2022-05-22 21:07:40 +02:00
|
|
|
|
right: Callable[[Node[T]], Optional[Node[T]]],
|
|
|
|
|
rotate_left: RotateFunction[T],
|
|
|
|
|
rotate_right: RotateFunction[T],
|
2022-05-20 00:01:04 +02:00
|
|
|
|
) -> Tuple[Optional[Node[T]], Optional[Node[T]]]:
|
2022-05-20 15:47:03 +02:00
|
|
|
|
if Node.difference(w) == Colour.Red:
|
2022-05-20 11:47:38 +02:00
|
|
|
|
# Case 1
|
|
|
|
|
# ======
|
|
|
|
|
# x’s sibling w is red
|
2022-05-22 21:00:18 +02:00
|
|
|
|
rotate_left(parent)
|
2022-05-20 00:01:04 +02:00
|
|
|
|
w = right(parent)
|
|
|
|
|
|
2022-05-20 15:47:03 +02:00
|
|
|
|
if Node.differences(w) == (Colour.Black, Colour.Black):
|
2022-05-20 11:47:38 +02:00
|
|
|
|
# Case 2
|
|
|
|
|
# ======
|
|
|
|
|
# x’s sibling w is black, and both of w’s children are black
|
2022-05-20 15:47:03 +02:00
|
|
|
|
parent.rank -= Colour.Black
|
2022-05-20 00:01:04 +02:00
|
|
|
|
x = parent
|
|
|
|
|
else:
|
2022-05-20 11:47:38 +02:00
|
|
|
|
# Case 3
|
|
|
|
|
# ======
|
|
|
|
|
# x’s sibling w is black,
|
|
|
|
|
# w’s left child is red, and w’s right child is black
|
2022-05-20 15:47:03 +02:00
|
|
|
|
if Node.difference(right(w), w) == Colour.Black:
|
2022-05-22 21:00:18 +02:00
|
|
|
|
rotate_right(w)
|
2022-05-20 00:01:04 +02:00
|
|
|
|
w = right(parent)
|
2022-05-20 11:47:38 +02:00
|
|
|
|
|
|
|
|
|
# Case 4
|
|
|
|
|
# ======
|
|
|
|
|
# x’s sibling w is black, and w’s right child is red
|
2022-05-20 15:47:03 +02:00
|
|
|
|
parent.rank -= Colour.Black
|
|
|
|
|
w.rank += Colour.Black
|
2022-05-22 21:00:18 +02:00
|
|
|
|
rotate_left(parent)
|
2022-05-20 00:01:04 +02:00
|
|
|
|
x = self.root
|
|
|
|
|
|
|
|
|
|
return x, (x.parent if x else None)
|
|
|
|
|
|
|
|
|
|
def _delete_rebalance(
|
|
|
|
|
self, node: Optional[Node[T]], parent: Optional[Node[T]]
|
|
|
|
|
) -> None:
|
|
|
|
|
if not node and not parent:
|
|
|
|
|
return
|
|
|
|
|
|
2022-05-22 00:38:48 +02:00
|
|
|
|
while node != self.root and has_double_black(parent):
|
2022-05-20 00:01:04 +02:00
|
|
|
|
if node == parent.left:
|
|
|
|
|
node, parent = self._delete_rebalance_step(
|
|
|
|
|
node,
|
|
|
|
|
parent.right,
|
|
|
|
|
parent,
|
|
|
|
|
lambda x: x.right,
|
2022-05-22 21:00:18 +02:00
|
|
|
|
self.rotate_left,
|
|
|
|
|
self.rotate_right,
|
2022-05-20 00:01:04 +02:00
|
|
|
|
)
|
|
|
|
|
else:
|
|
|
|
|
node, parent = self._delete_rebalance_step(
|
|
|
|
|
node,
|
|
|
|
|
parent.left,
|
|
|
|
|
parent,
|
|
|
|
|
lambda x: x.left,
|
2022-05-22 21:00:18 +02:00
|
|
|
|
self.rotate_right,
|
|
|
|
|
self.rotate_left,
|
2022-05-20 00:01:04 +02:00
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
# endregion DeleteRebalance
|