145 lines
3.8 KiB
Python
145 lines
3.8 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
|
|
|
|
def __fix_delete(
|
|
self,
|
|
x: Optional[Node[T]],
|
|
y: Node[T],
|
|
z: Node[T],
|
|
reversed: bool,
|
|
rotating_around_root: bool,
|
|
rotate_left: RotateFunction[T],
|
|
rotate_right: RotateFunction[T],
|
|
) -> None:
|
|
new_root = x
|
|
v = y.left
|
|
w = y.right
|
|
|
|
if reversed:
|
|
v, w = w, v
|
|
|
|
w_diff = Node.difference(w, y)
|
|
|
|
if w_diff == 1 and y.parent:
|
|
new_root = rotate_left(y.parent)
|
|
if rotating_around_root:
|
|
self.root = new_root
|
|
|
|
y.promote()
|
|
z.demote()
|
|
|
|
if z.type == NodeType.LEAF:
|
|
z.demote()
|
|
elif w_diff == 2 and v.parent:
|
|
rotate_right(v.parent)
|
|
new_root = rotate_left(v.parent)
|
|
if rotating_around_root:
|
|
self.root = new_root
|
|
|
|
v.promote().promote()
|
|
y.demote()
|
|
z.demote().demote()
|
|
|
|
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))
|
|
):
|
|
if y_diff != 2:
|
|
y.demote()
|
|
parent.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)
|
|
|
|
rotating_around_root = parent.parent is None
|
|
|
|
parent_node_diffs = Node.differences(parent)
|
|
if parent_node_diffs in ((1, 3), (3, 1)):
|
|
if parent.left is x:
|
|
assert parent.right
|
|
self.__fix_delete(
|
|
x,
|
|
parent.right,
|
|
parent,
|
|
False,
|
|
rotating_around_root,
|
|
Node.rotate_left,
|
|
Node.rotate_right,
|
|
)
|
|
else:
|
|
assert parent.left
|
|
self.__fix_delete(
|
|
x,
|
|
parent.left,
|
|
parent,
|
|
True,
|
|
rotating_around_root,
|
|
Node.rotate_right,
|
|
Node.rotate_left,
|
|
)
|
|
|
|
def _delete_rebalance(
|
|
self, y: Optional[Node[T]], parent: Optional[Node[T]]
|
|
) -> None:
|
|
if Node.differences(y) == (2, 2):
|
|
y.demote()
|
|
parent = y.parent
|
|
elif Node.differences(parent) == (2, 2):
|
|
parent.demote()
|
|
parent = parent.parent
|
|
|
|
if not parent:
|
|
return
|
|
|
|
for y in (parent.left, parent.right):
|
|
if Node.difference(y, parent) == 3:
|
|
self.__bottomup_delete(y, parent)
|
|
return
|
|
|
|
# endregion DeleteRebalance
|