summaryrefslogtreecommitdiffstats
path: root/utils.py
diff options
context:
space:
mode:
Diffstat (limited to 'utils.py')
-rw-r--r--utils.py60
1 files changed, 57 insertions, 3 deletions
diff --git a/utils.py b/utils.py
index 91b90ff..00d0d60 100644
--- a/utils.py
+++ b/utils.py
@@ -1,11 +1,65 @@
+from __future__ import annotations
from collections import deque
-from collections.abc import Iterable, Iterator, Generator
+from collections.abc import Iterable, Iterator
+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
+
+NumT = TypeVar('NumT', float, int)
+
+@dataclass(frozen=True)
+class Point2D(Generic[NumT]):
+ x: NumT
+ y: NumT
+
+ def __str__(self) -> str:
+ return f'({self.x}, {self.y})'
+
+ def __add__(self, other: Point2D[NumT]) -> Union[Point2D[NumT], bool]:
+ if isinstance(other, Point2D):
+ return Point2D(self.x + other.x, self.y + other.y)
+ return NotImplemented
+
+ def __sub__(self, other: Point2D[NumT]) -> Union[Point2D[NumT], bool]:
+ if isinstance(other, Point2D):
+ return Point2D(self.x - other.x, self.y - other.y)
+ return NotImplemented
+
+ def __abs__(self) -> Point2D[NumT]:
+ return Point2D(cast(NumT, abs(self.x)), cast(NumT, abs(self.y)))
+
+def norm_1(p: Point2D[NumT]) -> NumT:
+ return cast(NumT, abs(p.x) + abs(p.y))
+
+def norm_2(p: Point2D) -> float:
+ return sqrt(p.x ** 2 + p.y ** 2)
+
+def norm_inf(p: Point2D[NumT]) -> NumT:
+ return cast(NumT, max(abs(p.x), abs(p.y)))
+
+def inbounds(p: Point2D[NumT], a: Point2D[NumT], _b: Optional[Point2D[NumT]] = None) -> bool:
+ if isinstance(_b, Point2D):
+ b = _b
+ else:
+ b = a
+ a = cast(Point2D[NumT], Point2D(0, 0) if isinstance(a.x, int) else Point2D(0.0, 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) -> Iterator[Point2D[int]]:
+ 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) \
+ -> Iterator[Point2D[int]]:
+ 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]:
+def sliding_window(iterable: Iterable[T], n: int) -> Iterator[tuple[T, ...]]:
# sliding_window('ABCDEFG', 4) -> ABCD BCDE CDEF DEFG
it: Iterator[T] = iter(iterable)
window: deque[T] = deque(islice(it, n), maxlen=n)