python/node.py

from abc import abstractmethod
import enum
from typing import (
    Callable,
    Iterable,
    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})"

    def __iter__(self) -> Iterable[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 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]]