python/wavl.py

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