from node import Node, Comparable from abc import abstractmethod from collections import deque import logging from typing import Deque, 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) 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) y, parent = None, (n.parent if n.parent is not node else n) if n.parent is not node: parent = n.right if n.right else n.parent self._transplant(n, n.right) n.right = node.right n.right.parent = n self._transplant(node, n) n.left = node.left n.left.parent = 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)