151 lines
3.6 KiB
Python
151 lines
3.6 KiB
Python
from abc import abstractmethod
|
|
import enum
|
|
from typing import (
|
|
Iterable,
|
|
Iterator,
|
|
Optional,
|
|
Generic,
|
|
Tuple,
|
|
TypeVar,
|
|
Protocol,
|
|
)
|
|
|
|
|
|
class Comparable(Protocol):
|
|
"""Protocol for annotating comparable types."""
|
|
|
|
@abstractmethod
|
|
def __lt__(self: "Comparable", other: "Comparable") -> bool:
|
|
pass
|
|
|
|
|
|
class NodeType(enum.IntEnum):
|
|
LEAF = 0
|
|
UNARY = 1
|
|
BINARY = 2
|
|
|
|
|
|
T = TypeVar("T", bound=Comparable)
|
|
|
|
|
|
class Node(Generic[T], Iterable[T]):
|
|
def __init__(
|
|
self,
|
|
value: T,
|
|
left: "Optional[Node[T]]" = None,
|
|
right: "Optional[Node[T]]" = None,
|
|
parent: "Optional[Node[T]]" = None,
|
|
):
|
|
self.parent = parent
|
|
self.left = left
|
|
self.right = right
|
|
|
|
self.value = value
|
|
self.rank: int = 0
|
|
|
|
@property
|
|
def type(self) -> NodeType:
|
|
if self.left and self.right:
|
|
return NodeType.BINARY
|
|
|
|
if self.left or self.right:
|
|
return NodeType.UNARY
|
|
|
|
return NodeType.LEAF
|
|
|
|
def __repr__(self) -> str:
|
|
return (
|
|
f"Node(value={self.value}, rank={self.rank}, "
|
|
f"left={self.left}, right={self.right}, parent={self.parent})"
|
|
)
|
|
|
|
def __str__(self) -> str:
|
|
return f"Node(value={self.value}, rank={self.rank})"
|
|
|
|
def __iter__(self) -> Iterator[T]:
|
|
"""
|
|
Yields:
|
|
Keys from the subtree rooted at the node in an inorder fashion.
|
|
"""
|
|
if self.left:
|
|
yield from self.left
|
|
yield self.value
|
|
if self.right:
|
|
yield from self.right
|
|
|
|
@staticmethod
|
|
def height(node: "Optional[Node[T]]") -> int:
|
|
return (
|
|
1 + max(Node.height(node.left), Node.height(node.right))
|
|
if node
|
|
else -1
|
|
)
|
|
|
|
@staticmethod
|
|
def get_rank(node: "Optional[Node[T]]") -> int:
|
|
return -1 if not node else node.rank
|
|
|
|
@staticmethod
|
|
def difference(
|
|
node: "Optional[Node[T]]", parent: "Optional[Node[T]]" = None
|
|
) -> int:
|
|
if not parent:
|
|
parent = node.parent if node else None
|
|
return Node.get_rank(parent) - Node.get_rank(node)
|
|
|
|
@staticmethod
|
|
def differences(node: "Optional[Node[T]]") -> Tuple[int, int]:
|
|
node_rank = Node.get_rank(node)
|
|
(left, right) = (node.left, node.right) if node else (None, None)
|
|
|
|
return (
|
|
node_rank - Node.get_rank(left),
|
|
node_rank - Node.get_rank(right),
|
|
)
|
|
|
|
def promote(self) -> "Node[T]":
|
|
self.rank += 1
|
|
return self
|
|
|
|
def demote(self) -> "Node[T]":
|
|
self.rank -= 1
|
|
return self
|
|
|
|
@staticmethod
|
|
def find_parent_node(
|
|
value: T, node: "Node[T]", missing: bool = True
|
|
) -> "Optional[Node[T]]":
|
|
new_node: "Optional[Node[T]]" = node
|
|
|
|
while new_node and (missing or new_node.value != value):
|
|
node = new_node
|
|
if value < node.value:
|
|
new_node = node.left
|
|
elif node.value < value:
|
|
new_node = node.right
|
|
else:
|
|
return None
|
|
|
|
return node
|
|
|
|
@staticmethod
|
|
def search(value: T, node: "Optional[Node[T]]") -> "Optional[Node[T]]":
|
|
while node and node.value != value:
|
|
node = node.left if value < node.value else node.right
|
|
|
|
return node
|
|
|
|
@staticmethod
|
|
def minimum(node: "Node[T]") -> "Node[T]":
|
|
while node.left:
|
|
node = node.left
|
|
|
|
return node
|
|
|
|
@staticmethod
|
|
def maximum(node: "Node[T]") -> "Node[T]":
|
|
while node.right:
|
|
node = node.right
|
|
|
|
return node
|