142 lines
3.7 KiB
Python
142 lines
3.7 KiB
Python
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)
|
|
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)
|