python/ranked_tree.py

from node import Node, Comparable

from abc import abstractmethod
from collections import deque
import logging
from typing import Deque, Iterable, Optional, Tuple, TypeVar, Generic

logger = logging.getLogger(__name__)
T = TypeVar("T", bound=Comparable)


class RankedTree(Generic[T]):
    def __init__(self) -> None:
        self.root: Optional[Node[T]] = None

    def _get_unwrapped_graph(self) -> str:
        result = ""

        queue: Deque[Optional[Node[T]]] = deque([self.root])
        edges = []
        while queue:
            node = queue.popleft()
            if not node:
                continue

            result += f'"{str(node)}" [label="{node.value}, {node.rank}"];\n'
            for child in (node.left, node.right):
                if not child:
                    continue

                edges.append((node, child))
                queue.append(child)

        for from_node, to_node in edges:
            label = f'[label="{Node.difference(to_node)}"]'
            result += f'"{str(from_node)}" -> "{str(to_node)}" {label}\n'

        return result

    def __str__(self) -> str:
        return "digraph {\n" + self._get_unwrapped_graph() + "}\n"

    # region TreeSpecificMethods

    @abstractmethod
    def is_correct_node(self, node: Optional[Node[T]]) -> bool:
        pass

    @abstractmethod
    def _insert_rebalance(self, x: Node[T]) -> None:
        pass

    @abstractmethod
    def _delete_rebalance(
        self, node: Optional[Node[T]], parent: Optional[Node[T]]
    ) -> None:
        pass

    # endregion TreeSpecificMethods

    @property
    def rank(self) -> int:
        return Node.get_rank(self.root)

    @property
    def is_correct(self) -> bool:
        return self.is_correct_node(self.root)

    def search(
        self, value: T, node: Optional[Node[T]] = None
    ) -> Optional[Node[T]]:
        if not node:
            node = self.root

        return Node.search(value, node)

    def insert(self, value: T) -> None:
        inserted_node = Node(value)

        if not self.root:
            self.root = inserted_node
            return

        parent = Node.find_parent_node(value, self.root)
        if not parent:
            return

        inserted_node.parent = parent

        if value < parent.value:
            parent.left = inserted_node
        else:
            parent.right = inserted_node

        self._insert_rebalance(inserted_node)

    def _transplant(self, u: Node[T], v: Optional[Node[T]]) -> None:
        if not u.parent:
            self.root = v
        elif u.parent.left is u:
            u.parent.left = v
        else:
            u.parent.right = v

        if v:
            v.rank = u.rank
            v.parent = u.parent

    def _delete_node(
        self, node: Optional[Node[T]]
    ) -> Optional[Tuple[Optional[Node[T]], Optional[Node[T]]]]:
        if node is None:
            return None

        y, parent = None, node.parent

        if not node.left:
            y = node.right
            self._transplant(node, node.right)
        elif not node.right:
            y = node.left
            self._transplant(node, node.left)
        else:
            n = Node.minimum(node.right)
            node.value = n.value
            n.value = None
            return self._delete_node(n)

        return (y, parent)

    def _delete(
        self, value: T
    ) -> Optional[Tuple[Optional[Node[T]], Optional[Node[T]]]]:
        node = self.root
        while node is not None and node.value != value:
            node = node.left if value < node.value else node.right
        return self._delete_node(node)

    def delete(self, value: T) -> None:
        if to_be_rebalanced := self._delete(value):
            y, parent = to_be_rebalanced
            self._delete_rebalance(y, parent)

    def __iter__(self) -> Iterable[T]:
        """
        Yields:
            Keys from the tree in an inorder fashion.
        """
        if self.root:
            yield from self.root