python/wavl.py
Matej Focko 4a832e1de8
feat: add avl and split wavl
Signed-off-by: Matej Focko <mfocko@redhat.com>
2022-01-30 15:03:34 +01:00

165 lines
4.5 KiB
Python

from avl import AVLTree
from node import Node, NodeType, Comparable, RotateFunction
import logging
from typing import TypeVar, Optional
logger = logging.getLogger(__name__)
T = TypeVar("T", bound=Comparable)
class WAVLTree(AVLTree[T]):
def is_correct_node(
self, node: Optional[Node[T]], recursive: bool = True
) -> bool:
if not node:
return True
for child_rank in Node.differences(node):
if child_rank not in (1, 2):
return False
if node.type == NodeType.LEAF:
return node.rank == 0
return not recursive or (
self.is_correct_node(node.left)
and self.is_correct_node(node.right)
)
# region DeleteRebalance
@staticmethod
def __fix_delete(
x: Optional[Node[T]],
y: Node[T],
z: Node[T],
reversed: bool,
rotate_left: RotateFunction[T],
rotate_right: RotateFunction[T],
) -> Optional[Node[T]]:
new_root = x
v = y.left
w = y.right
if reversed:
v, w = w, v
logger.debug(f"__fix_delete({x}, {y}, {z}, {reversed})")
w_diff = Node.difference(w, y)
logger.debug(f"w_diff = {w_diff}")
assert v
if w_diff == 1 and y.parent:
logger.debug(f"y.parent = {y.parent}")
new_root = rotate_left(y.parent)
y.promote()
z.demote()
if z.type == NodeType.LEAF:
z.demote()
elif w_diff == 2 and v.parent:
logger.debug(f"v.parent = {v.parent}")
rotate_right(v.parent)
new_root = rotate_left(v.parent)
v.promote().promote()
y.demote()
z.demote().demote()
return new_root
def __bottomup_delete(
self, x: Optional[Node[T]], parent: Optional[Node[T]]
) -> None:
x_diff = Node.difference(x, parent)
if x_diff != 3 or not parent:
return
y = parent.right if parent.left is x else parent.left
y_diff = Node.difference(y, parent)
while (
parent
and x_diff == 3
and y
and (y_diff == 2 or Node.differences(y) == (2, 2))
):
parent.demote()
if y_diff != 2:
y.demote()
x = parent
parent = x.parent
if not parent:
return
y = parent.right if parent.left is x else parent.left
x_diff = Node.difference(x, parent)
y_diff = Node.difference(y, parent)
if not parent:
return
rotating_around_root = parent.parent is None
new_root: Optional[Node[T]] = parent
parent_node_diffs = Node.differences(parent)
if parent_node_diffs in ((1, 3), (3, 1)):
if parent.left is x:
assert parent.right
new_root = WAVLTree.__fix_delete(
x,
parent.right,
parent,
False,
Node.rotate_left,
Node.rotate_right,
)
else:
assert parent.left
new_root = WAVLTree.__fix_delete(
x,
parent.left,
parent,
True,
Node.rotate_right,
Node.rotate_left,
)
if rotating_around_root:
self.root = new_root
def __delete_fixup(
self, y: Optional[Node[T]], parent: Optional[Node[T]] = None
) -> None:
logger.debug(f"[__delete_fixup] y = {y}, parent = {parent}")
z = y if y else parent
logger.debug(
f"[z.demote()] Node.differences({repr(z)}) == (2, 2) ~>*"
f"{Node.differences(z)} == (2, 2)"
)
assert z
if Node.differences(z) == (2, 2):
z.demote()
if parent:
for y in (parent.left, parent.right):
logger.debug(
f"[bottom-up delete] Node.difference({y}, {parent}) == 3"
f"~>* {Node.difference(y, parent)} == 3"
)
if Node.difference(y, parent) == 3:
self.__bottomup_delete(y, parent)
def _delete_rebalance(
self, node: Optional[Node[T]], parent: Optional[Node[T]]
) -> None:
while node or parent:
self.__delete_fixup(node, parent)
node, parent = parent, (parent.parent if parent else None)
# endregion DeleteRebalance