Some helpers function for competitive programming

utils.py

"""LuckyJ.: Utils scripts for CP"""

from __future__ import annotations

import subprocess
import string

from heapq import heappop, heappush
from queue import Queue
from queue import LifoQueue as Stack
from typing import Callable, Dict, Iterator, List
from typing import Sequence, Set, Tuple, TypeVar, Union

#######
#  GLOBAL VARIABLES
#####

DX8 = [1, -1, 0, 0, 1, -1, -1, 1]
DY8 = [0, 0, 1, -1, 1, -1, 1, -1]
DX4 = [1, -1, 0, 0]
DY4 = [0, 0, 1, -1]
DIR = ["R", "L", "U", "D", "RU", "LD", "LU", "RD"]

INF = float("inf")
cat = "".join

T = TypeVar("T")

#######
#  DECORATORS
#####


def static_vars(**kwargs):
    def decorate(func):
        for k in kwargs:
            setattr(func, k, kwargs[k])
        return func

    return decorate


#######
#  CONVERTERS
#####


def ints(s: Union[str, List[T]], sep=None) -> List[int]:
    """Converts to a list of int:
    - a string
    - an iterable
    """
    if isinstance(s, str) and not sep is None:
        return [int(x) for x in s.split(sep)]
    else:
        return [int(x) for x in s]


def flatten(arr: List[List[T]]) -> List[T]:
    return [val for row in arr for val in row]


@static_vars(digs=string.digits + string.ascii_letters)
def int2base(x: int, base: int) -> List[str]:
    """Convert an int (base 10) to any base"""
    if x < 0:
        sign = -1
    elif x == 0:
        return int2base.digs[0]
    else:
        sign = 1

    x *= sign
    digits = []

    while x:
        digits.append(int2base.digs[x % base])
        x = x // base

    if sign < 0:
        digits.append("-")

    return digits[::-1]


#######
#  PRINTERS
#####


def cpp(s: T) -> None:
    """Print `s` and do a copy of it to the cleapboard"""
    print(s)
    subprocess.Popen(
        ["xclip", "-selection", "clipboard"], stdin=subprocess.PIPE
    ).communicate(bytes(str(s), encoding="ascii"))


def print_grid(grid: List[List[T]], align: str = None) -> None:
    """Print a grid"""
    if align is not None:
        assert align in list("><^lrm"), f'"{align}" not in "><^lrm"'
        if align in "lrm":
            align = "<>^"["lrm".index(align)]
        frmt_col = [
            f"{align}{max(map(lambda x: len(str(x)), row))}" for row in transpose(grid)
        ]
        for row in grid:
            for col, frmt in zip(row, frmt_col):
                print(f"{col:{frmt}}", end=" ")
            print()
    else:
        for row in grid:
            print(*row)


#######
#  MATRICES
#####


def transpose(matrix: List[List[T]]) -> List[List[T]]:
    """Transpose a matrix"""
    return [*zip(*matrix)]


def isValid(
    x: int, y: int, max_x: int = None, max_y: int = None, min_x: int = 0, min_y: int = 0
) -> bool:
    """Check if min_x <= x <= max_x and min_y <= y <= max_y"""
    return (
        min_x <= x <= (max_x, x)[max_x is None]
        and min_y <= y <= (max_y, y)[max_y is None]
    )


def neighbors4(
    x: int, y: int, max_x: int = None, max_y: int = None, min_x: int = 0, min_y: int = 0
) -> List[Tuple[int, int]]:
    return [
        (x + dx, y + dy)
        for dx, dy in zip(DX4, DY4)
        if isValid(x + dx, y + dy, max_x, max_y, min_x, min_y)
    ]


def neighbors8(
    x: int, y: int, max_x: int = None, max_y: int = None, min_x: int = 0, min_y: int = 0
) -> List[Tuple[int, int]]:
    return [
        (x + dx, y + dy)
        for dx, dy in zip(DX8, DY8)
        if isValid(x + dx, y + dy, max_x, max_y, min_x, min_y)
    ]


#######
#  SEARCH
#####


def search(arr: List[T], val: T) -> int:
    """Performs binary search on an array"""
    l, r = 0, len(arr) - 1
    while l <= r:
        m = l + (r - l) // 2
        if arr[m] == val:
            return m
        if val < arr[m]:
            r = m - 1
        else:
            l = m + 1
    return -1


def _path(previous: Dict[T, T], s: T) -> List[T]:
    """Return a list of states that lead to state s, according to the previous dict."""
    return [] if (s is None) else _path(previous, previous[s]) + [s]


def euclidean(x_1: int, y_1: int, x_2: int, y_2: int) -> float:
    """Euclidean distance between two points"""
    return (abs(x_2 - x_1) ** 2 + abs(y_2 - y_1) ** 2) ** 0.5


def astar_search(
    start_x: int,
    start_y: int,
    end_x: int,
    end_y: int,
    h_func: Callable[[int, int, int, int], float] = euclidean,
    move_func: Callable[[int, int], List[Tuple[int, int]]] = neighbors4,
) -> Tuple[bool, List[T]]:
    """Find a shortest sequence of states from start to end."""
    start = (start_x, start_y)
    end = (end_x, end_y)
    hq: List[Tuple[float, Tuple[int, int]]] = [(h_func(*start, *end), start)]
    prev: Dict[Tuple[int, int], Tuple[int, int]] = {start: None}
    cost: Dict[Tuple[int, int], float] = {start: 0.0}
    while hq:
        _, curr = heappop(hq)
        if h_func(*curr, *end) == 0:
            return True, _path(prev, curr)
        for next in move_func(*curr):
            new_cost = cost[curr] + 1
            if next not in cost or new_cost < cost[next]:
                heappush(hq, (new_cost + h_func(*next, *end), next))
                cost[next] = new_cost
                prev[next] = curr
    return False, []


def BFS(graph: Dict[T, List[T]], start: T, end: T):
    if start == end:
        return True

    visited: Set[T] = set()
    queue: Queue = Queue()

    visited.add(start)
    queue.put(start)
    while not queue.empty():
        curr = queue.get()
        for next in graph[curr]:
            if next == end:
                return True
            if not next in visited:
                visited.add(next)
                queue.put(next)
    return False


def DFS(graph: Dict[T, List[T]], start: T, end: T):
    if start == end:
        return True

    visited: Set[T] = set()
    stack: Stack = Stack()

    visited.add(start)
    stack.put(start)
    while not stack.empty():
        curr = stack.get()
        for next in graph[curr]:
            if next == end:
                return True
            if not next in visited:
                visited.add(next)
                stack.put(next)
    return False


#######
#  MISC ALGORITHMS
#####


def GCD(x: int, y: int) -> int:
    return x if y == 0 else GCD(y, x % y)


def dijkstra(graph: Dict[T, Dict[T, float]], start: T) -> Dict[T, float]:
    """Apply Dijkstra and get the Dict of dist to a start point"""
    dijk: Dict[T, float] = dict()
    hq: List[Tuple[float, T]] = list()

    for node in graph:
        dijk[node] = INF
    dijk[start] = 0

    heappush(hq, (dijk[start], start))
    while hq:
        dist, curr = heappop(hq)
        dijk[curr] = min(dijk[curr], dist)
        for next in graph[curr]:
            if dijk[next] > dijk[curr] + graph[curr][next]:
                dijk[next] = dijk[curr] + graph[curr][next]
                heappush(hq, (dijk[next], next))
    return dijk


def LCS(s1: str, s2: str) -> str:
    """Answer the "longest common subsequence" problem
    - Runtime: O(len(s1) * len(s2))
    - Space  : O(len(s1) * len(s2))
    """
    dp = [[0 for _ in range(len(s2) + 1)] for _ in range(len(s1) + 1)]
    for i, c1 in enumerate(s1):
        for j, c2 in enumerate(s2):
            if c1 == c2:
                dp[i + 1][j + 1] = dp[i][j] + 1
            else:
                dp[i + 1][j + 1] = max(dp[i][j + 1], dp[i + 1][j])

    lcs = ""
    i, j = len(s1), len(s2)
    while i >= 1 and j >= 1:
        if s1[i - 1] == s2[j - 1]:
            lcs += s1[i - 1]
            i, j = i - 1, j - 1
        elif dp[i - 1][j] > dp[i][j - 1]:
            i -= 1
        else:
            j -= 1
    return lcs[::-1]


#######
#  VECTOR
#####


class Vector(object):
    def __init__(self, *vals, dtype=None):
        self.__dtype = type(vals[0]) if dtype is None else dtype
        if len(vals) == 1:
            self.__vec = list(vals[0])
        else:
            self.__vec = list(map(self.__dtype, vals))

    ### GETTERS / SETTERS ###

    @property
    def x(self):
        return self.__vec[0]

    @property
    def y(self):
        return self.__vec[1]

    @property
    def z(self):
        return self.__vec[2]

    @x.setter
    def x(self, val):
        self.__vec[0] = self.__dtype(val)

    @y.setter
    def y(self, val):
        self.__vec[1] = self.__dtype(val)

    @z.setter
    def z(self, val):
        self.__vec[2] = self.__dtype(val)

    def __getitem__(self, idx: int):
        return self.__vec[idx]

    def __setitem__(self, idx: int, val) -> None:
        self.__vec[idx] = self.__dtype(val)

    ### SEQUENCE ###

    def __len__(self) -> int:
        return len(self.__vec)

    def __iter__(self) -> Iterator:
        return iter(self.__vec)

    def append(self, val) -> None:
        """Adds an element at the end of the vector."""
        self.__vec.append(self.__dtype(val))

    def clear(self) -> None:
        """Removes all the elements from the vector."""
        self.__vec.clear()

    def copy(self, dtype=None) -> Vector:
        """Returns a copy of the vector."""
        return Vector(*self.__vec.copy(), self.__dtype if dtype is None else dtype)

    def count(self, val) -> int:
        """Returns the number of elements with the specified value."""
        return self.__vec.count(self.__dtype(val))

    def extend(self, other) -> None:
        """Add the elements of a vector (or any iterable), to the end of the current vector."""
        self.__vec.extend([self.__dtype(val) for val in other])

    def index(self, val) -> int:
        """Returns the index of the first element with the specified value."""
        self.__vec.index(self.__dtype(val))

    def insert(self, idx: int, val) -> None:
        """."""
        self.__vec.insert(idx, self.__dtype(val))

    def pop(self, idx: int) -> T:
        """Removes the element at the specified position."""
        return self.__vec.pop(idx)

    def remove(self, val) -> None:
        """Removes the first item with the specified value."""
        self.__vec.remove(self.__dtype(val))

    def reverse(self) -> None:
        """Reverses the order of the vector."""
        return self.__vec.reverse()

    def sort(self) -> None:
        """Sorts the vector."""
        return self.__vec.sort()

    ### + - * // / == != ###

    def __add__(self, other: Union[float, Sequence[float]]) -> Vector:
        """`self + other` Element-wise add."""
        try:
            return Vector(*(a + b for a, b in zip(self, other)))
        except:
            return Vector(*(a + other for a in self))

    def __sub__(self, other: Union[float, Sequence[float]]) -> Vector:
        """`self - other` Element-wise sub. Strict trunc to smallest Sequence."""
        try:
            return Vector(*(a - b for a, b in zip(self, other)))
        except:
            return Vector(*(a - other for a in self))

    def __mul__(self, other: Union[float, Sequence[float]]) -> Vector:
        """`self * other` Element-wise mul. Strict trunc to smallest Sequence."""
        try:
            return Vector(*(a * b for a, b in zip(self, other)))
        except:
            return Vector(*(a * other for a in self))

    def __truediv__(self, other: Union[float, Sequence[float]]) -> Vector:
        """`self / other` Element-wise mul. Strict trunc to smallest Sequence."""
        try:
            return Vector(*(a / b for a, b in zip(self, other)))
        except:
            return Vector(*(a / other for a in self))

    def __floordiv__(self, other: Union[float, Sequence[float]]) -> Vector:
        """`self // other` Element-wise mul. Strict trunc to smallest Sequence."""
        try:
            return Vector(*(a // b for a, b in zip(self, other)))
        except:
            return Vector(*(a // other for a in self))

    def __radd__(self, other: Union[float, Sequence[float]]) -> Vector:
        """`other + self`"""
        return self.__add__(other)

    def __rmul__(self, other: Union[float, Sequence[float]]) -> Vector:
        """`other * self`"""
        return self.__mul__(other)

    def __iadd__(self, other: Union[float, Sequence[float]]) -> Vector:
        """`self += other`"""
        self = self.__add__(other)
        return self

    def __isub__(self, other: Union[float, Sequence[float]]) -> Vector:
        """`self -= other`"""
        self = self.__sub__(other)
        return self

    def __imul__(self, other: Union[float, Sequence[float]]) -> Vector:
        """`self *= other`"""
        self = self.__mul__(other)
        return self

    def __itruediv__(self, other: Union[float, Sequence[float]]) -> Vector:
        """`self /= other`"""
        self = self.__truediv__(other)
        return self

    def __ifloordiv__(self, other: Union[float, Sequence[float]]) -> Vector:
        """`self //= other`"""
        self = self.__floordiv__(other)
        return self

    def __neg__(self) -> Vector:
        """`-self`"""
        return Vector(*(-a for a in self))

    def __eq__(self, other):
        return tuple(self) == tuple(other)

    def __ne__(self, other):
        return not self.__eq__(other)

    ### MISCELLANEOUS ###

    def __str__(self) -> str:
        return "<{}>".format(", ".join(map(str, self)))

    def __repr__(self) -> str:
        return str(self)

    def __hash__(self):
        return hash(tuple(self))

    def norm(self):
        return sum(a * a for a in self) ** 0.5

    @property
    def normalized(self):
        return self / self.norm()