python/ranked_tree.py
Matej Focko 4a832e1de8
feat: add avl and split wavl
Signed-off-by: Matej Focko <mfocko@redhat.com>
2022-01-30 15:03:34 +01:00

146 lines
3.9 KiB
Python

from abc import abstractmethod
from node import Node, Comparable
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 __str__(self) -> str:
result = "digraph {\n"
queue: Deque[Optional[Node[T]]] = deque()
queue.append(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 + "}\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)