diff options
Diffstat (limited to 'utils.py')
| -rw-r--r-- | utils.py | 58 | 
1 files changed, 57 insertions, 1 deletions
| @@ -1,8 +1,64 @@  from collections import deque  from collections.abc import Iterable, Iterator, Generator +from dataclasses import dataclass  from itertools import islice +from math import sqrt  from sys import argv -from typing import TypeVar +from typing import TypeVar, Generic, Union, Optional, cast + +Numeric = Union[int, float] + +T = TypeVar('T', bound=Numeric) +@dataclass(frozen=True) +class Point2D(Generic[T]): +    x: T +    y: T + +    def __str__(self) -> str: +        return f'({self.x}, {self.y})' + +    def __add__(self, other: 'Point2D') -> Union['Point2D', bool]: +        if isinstance(other, Point2D): +            return Point2D(self.x + other.x, self.y + other.y) +        return NotImplemented + +    def __sub__(self, other: 'Point2D') -> Union['Point2D', bool]: +        if isinstance(other, Point2D): +            return Point2D(self.x - other.x, self.y - other.y) +        return NotImplemented + +    def __abs__(self) -> 'Point2D': +        return Point2D(abs(self.x), abs(self.y)) + +T = TypeVar('T', bound=Numeric) +def norm_1(p: Point2D[T]) -> T: +    return abs(p.x) + abs(p.y) + +def norm_2(p: Point2D) -> float: +    return sqrt(p.x ** 2 + p.y ** 2) + +T = TypeVar('T', bound=Numeric) +def norm_inf(p: Point2D[T]) -> T: +    return max(abs(p.x), abs(p.y)) + +T = TypeVar('T', bound=Numeric) +def inbounds(p: Point2D[T], a: Point2D[T], b: Optional[Point2D[T]] = None) -> bool: +    if b is None and isinstance(a.x, int): +        b = a +        a = cast(Point2D[T], Point2D(0, 0)) +    return p.x >= a.x and p.y >= a.y and p.x < b.x and p.y < b.y + +def adjacent(p: Point2D[int], diagonal: bool = True) \ +        -> Generator[Point2D[int], None, None]: +    for dx in range(-1, 2): +        for dy in range(-1, 2): +            if dx == 0 and dy == 0: continue +            if dx != 0 and dy != 0 and not adjacent: continue +            yield Point2D(p.x + dx, p.y + dy) + +def adjacent_bounded(p: Point2D[int], bound: Point2D[int], diagonal: bool = True) \ +        -> filter[Point2D[int]]: # This has to return filter[...] not Generator[...] because ??? +    return filter(lambda p: inbounds(p, bound), adjacent(p, diagonal))  T = TypeVar('T')  def sliding_window(iterable: Iterable[T], n: int) -> Generator[tuple[T, ...], None, None]: | 
