180 lines
4.2 KiB
Python
180 lines
4.2 KiB
Python
from abc import abstractmethod
|
|
import enum
|
|
from typing import Callable, 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]):
|
|
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})"
|
|
|
|
@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 rotate_right(x: "Node[T]") -> "Node[T]":
|
|
parent = x.parent
|
|
y = x.left
|
|
# z = x.right
|
|
|
|
assert y is not None
|
|
if parent:
|
|
if parent.left is x:
|
|
parent.left = y
|
|
else:
|
|
parent.right = y
|
|
|
|
x.left = y.right
|
|
if x.left:
|
|
x.left.parent = x
|
|
|
|
y.right = x
|
|
x.parent = y
|
|
y.parent = parent
|
|
|
|
return y
|
|
|
|
@staticmethod
|
|
def rotate_left(x: "Node[T]") -> "Node[T]":
|
|
parent = x.parent
|
|
# y = x.left
|
|
z = x.right
|
|
|
|
assert z is not None
|
|
if parent:
|
|
if parent.left is x:
|
|
parent.left = z
|
|
else:
|
|
parent.right = z
|
|
|
|
x.right = z.left
|
|
if x.right:
|
|
x.right.parent = x
|
|
|
|
z.left = x
|
|
x.parent = z
|
|
z.parent = parent
|
|
|
|
return z
|
|
|
|
@staticmethod
|
|
def find_parent_node(
|
|
value: T, node: "Node[T]", missing: bool = True
|
|
) -> "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
|
|
|
|
|
|
RotateFunction = Callable[[Node[T]], Node[T]]
|