165 lines
4.4 KiB
Python
165 lines
4.4 KiB
Python
from node import Node, Comparable, RotateFunction
|
|
from ranked_tree import RankedTree
|
|
|
|
import logging
|
|
from typing import TypeVar, Optional
|
|
|
|
logger = logging.getLogger(__name__)
|
|
T = TypeVar("T", bound=Comparable)
|
|
|
|
|
|
def _balance_factor(node: Optional[Node[T]]) -> int:
|
|
if not node:
|
|
return 0
|
|
|
|
left, right = Node.get_rank(node.left), Node.get_rank(node.right)
|
|
return right - left
|
|
|
|
|
|
def _update_rank(node: Node[T]) -> None:
|
|
left, right = Node.get_rank(node.left), Node.get_rank(node.right)
|
|
# logger.debug(f"[_update_rank] on {node} = ({left}, {right})")
|
|
node.rank = 1 + max(left, right)
|
|
|
|
|
|
class AVLTree(RankedTree[T]):
|
|
def is_correct_node(
|
|
self, node: Optional[Node[T]], recursive: bool = True
|
|
) -> bool:
|
|
if not node:
|
|
return True
|
|
|
|
if not (-1 <= _balance_factor(node) <= 1):
|
|
return False
|
|
|
|
return not recursive or (
|
|
self.is_correct_node(node.left)
|
|
and self.is_correct_node(node.right)
|
|
)
|
|
|
|
# region InsertRebalance
|
|
|
|
def __fix_0_child(
|
|
self,
|
|
x: Node[T],
|
|
y: Optional[Node[T]],
|
|
z: Node[T],
|
|
rotate_left: RotateFunction[T],
|
|
rotate_right: RotateFunction[T],
|
|
) -> Optional[Node[T]]:
|
|
new_root = x.parent
|
|
|
|
if not y or Node.difference(y) == 2:
|
|
new_root = rotate_right(z)
|
|
z.demote()
|
|
elif Node.difference(y) == 1:
|
|
rotate_left(x)
|
|
new_root = rotate_right(z)
|
|
|
|
y.promote()
|
|
x.demote()
|
|
z.demote()
|
|
|
|
return new_root
|
|
|
|
def _insert_rebalance(self, x: Node[T]) -> None:
|
|
diffs = Node.differences(x.parent)
|
|
while x.parent and (diffs == (0, 1) or diffs == (1, 0)):
|
|
x.parent.promote()
|
|
x = x.parent
|
|
|
|
diffs = Node.differences(x.parent)
|
|
|
|
if not x.parent:
|
|
return
|
|
|
|
rotating_around_root = x.parent.parent is None
|
|
new_root: Optional[Node[T]] = x.parent
|
|
|
|
rank_difference = Node.difference(x)
|
|
if rank_difference != 0:
|
|
return
|
|
|
|
x_parent = x.parent
|
|
assert x_parent is not None
|
|
if rank_difference == 0 and x.parent.left is x:
|
|
new_root = self.__fix_0_child(
|
|
x, x.right, x_parent, Node.rotate_left, Node.rotate_right
|
|
)
|
|
elif rank_difference == 0 and x.parent.right is x:
|
|
new_root = self.__fix_0_child(
|
|
x, x.left, x_parent, Node.rotate_right, Node.rotate_left
|
|
)
|
|
|
|
if rotating_around_root:
|
|
self.root = new_root
|
|
|
|
# endregion InsertRebalance
|
|
|
|
# region DeleteRebalance
|
|
|
|
def __delete_rotate(
|
|
self,
|
|
x: Node[T],
|
|
y: Node[T],
|
|
leaning: int,
|
|
rotating_around_root: bool,
|
|
rotate_left: RotateFunction[T],
|
|
rotate_right: RotateFunction[T],
|
|
) -> bool:
|
|
new_root = x
|
|
|
|
factor = _balance_factor(y)
|
|
if factor in (0, leaning):
|
|
new_root = rotate_left(x)
|
|
else:
|
|
rotate_right(y)
|
|
new_root = rotate_left(x)
|
|
|
|
for n in filter(None, (new_root.left, new_root.right, new_root)):
|
|
_update_rank(n)
|
|
|
|
if rotating_around_root:
|
|
self.root = new_root
|
|
return factor == 0
|
|
|
|
def __delete_fixup(
|
|
self, y: Optional[Node[T]], parent: Optional[Node[T]] = None
|
|
) -> bool:
|
|
x = y if y else parent
|
|
assert x
|
|
|
|
factor = _balance_factor(x)
|
|
if factor == 0:
|
|
_update_rank(x)
|
|
return False
|
|
elif factor in (-1, 1):
|
|
return True
|
|
|
|
rotating_around_root = x.parent is None
|
|
y, leaning, to_left, to_right = (
|
|
(x.right, 1, Node.rotate_left, Node.rotate_right)
|
|
if factor == 2
|
|
else (x.left, -1, Node.rotate_right, Node.rotate_left)
|
|
)
|
|
assert y
|
|
return self.__delete_rotate(
|
|
x,
|
|
y,
|
|
leaning,
|
|
rotating_around_root,
|
|
to_left,
|
|
to_right,
|
|
)
|
|
|
|
def _delete_rebalance(
|
|
self, node: Optional[Node[T]], parent: Optional[Node[T]]
|
|
) -> None:
|
|
while node or parent:
|
|
# TODO: Check if it is possible to not propagate all the way up.
|
|
# if self.__delete_fixup(node, parent):
|
|
# return
|
|
self.__delete_fixup(node, parent)
|
|
node, parent = parent, (parent.parent if parent else None)
|
|
|
|
# endregion DeleteRebalance
|