python/node.py
Matej Focko bc981a8bb4
node: add type signatures
Signed-off-by: Matej Focko <mfocko@redhat.com>
2022-01-30 15:02:28 +01:00

175 lines
4.1 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
new_node = node.left if value < node.value else node.right
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]]