box-o-sand/leetcode/stdlib.py

import dataclasses
import typing


class LinkedListNode(typing.Protocol):
    val: int
    next: typing.Optional["LinkedListNode"]


class ListNode:
    """ListNode is the leetcode "standard library" type used in linked lists"""

    def __init__(self, val: int = 0, next: typing.Optional["ListNode"] = None):  # no qa
        self.val = val
        self.next = next


class ListNodeRandom:
    """ListNodeRandom is another weirdo linked list thing from
    leetcode that is obviously very different than a binary tree
    node :upside_down_face:
    """

    def __init__(
        self,
        x: int,
        next: typing.Optional["ListNodeRandom"] = None,
        random: typing.Optional["ListNodeRandom"] = None,
    ):
        self.val = x
        self.next = next
        self.random = random


@dataclasses.dataclass
class ListNodeRandomNicely:
    val: int
    next: typing.Optional["ListNodeRandomNicely"] = None
    random: typing.Optional["ListNodeRandomNicely"] = None


class BinaryTreeNode(typing.Protocol):
    val: int
    left: typing.Optional["BinaryTreeNode"]
    right: typing.Optional["BinaryTreeNode"]


class TreeNode:
    """TreeNode is the leetcode "standard library" type used in binary trees"""

    def __init__(
        self,
        val: int = 0,
        left: typing.Optional["TreeNode"] = None,
        right: typing.Optional["TreeNode"] = None,
    ):
        self.val = val
        self.left = left
        self.right = right

    @classmethod
    def from_int(cls, val: int | None) -> typing.Optional["TreeNode"]:
        if val is None:
            return None

        return TreeNode(val)

    # __repr__ was added by me
    def __repr__(self) -> str:
        filtered_parts = []

        for key, value in [
            ("val", self.val),
            ("left", self.left),
            ("right", self.right),
        ]:
            if value is not None:
                filtered_parts.append((key, value))

        middle = ", ".join([f"{k}={v!r}" for k, v in filtered_parts])

        return f"TreeNode({middle})"

    # __eq__ was added by me
    def __eq__(self, other: typing.Optional["TreeNode"]) -> bool:
        return (
            other is not None
            and self.val == other.val
            and self.left == other.left
            and self.right == other.right
        )


class ConnectableBinaryTreeNode(typing.Protocol):
    val: int
    left: typing.Optional["BinaryTreeNode"]
    right: typing.Optional["BinaryTreeNode"]
    next: typing.Optional["BinaryTreeNode"]


class Node:
    """Node is the *other* leetcode "standard library" type used in binary trees"""

    def __init__(
        self,
        val: int = 0,
        left: typing.Optional["Node"] = None,
        right: typing.Optional["Node"] = None,
        next: typing.Optional["Node"] = None,  # no qa
    ):
        self.val = val
        self.left = left
        self.right = right
        self.next = next

    # __repr__ was added by me
    def __repr__(self) -> str:
        filtered_parts = []

        for key, value in [
            ("val", self.val),
            ("right", self.right),
            ("left", self.left),
            ("next", self.next),
        ]:
            if value is not None:
                filtered_parts.append((key, value))

        middle = ", ".join([f"{k}={v!r}" for k, v in filtered_parts])

        return f"Node({middle})"

    # __eq__ was added by me
    def __eq__(self, other: "Node") -> bool:
        return (
            other is not None
            and self.val == other.val
            and self.left == other.left
            and self.right == other.right
            and self.next == other.next
        )

    # __list__ was added by me
    def __list__(self) -> list[int | None]:
        ret = [self.val]
        ret += self.right.__list__() if self.right is not None else [None]
        ret += self.left.__list__() if self.left is not None else [None]
        ret += self.next.__list__() if self.next is not None else [None]
        return ret


class NeighborlyNode:
    """NeighborlyNode is a "Node" type used in leetcode graph puzzles"""

    def __init__(self, val=0, neighbors=None):
        self.val = val
        self.neighbors = neighbors if neighbors is not None else []


class NeighborlyNodeNicely(typing.NamedTuple):
    val: int
    neighbors: list["NeighborlyNodeNicely"]

    def __eq__(self, other: typing.Optional["NeighborlyNodeNicely"]) -> bool:
        return (
            other is not None
            and self.val == other.val
            and [n.val for n in self.neighbors] == [n.val for n in other.neighbors]
        )