from __future__ import annotations from collections import deque, defaultdict from collections.abc import Iterable, Iterator from dataclasses import dataclass from heapq import heappop, heappush from itertools import islice from math import sqrt from re import match as re_match from sys import argv from typing import TypeVar, Generic, Union, Optional, cast NumT = TypeVar('NumT', float, int) @dataclass(frozen=True, slots=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: b: Point2D[NumT] 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 diagonal: 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)) def a_star(start, goal, neighbours, h, d): vtoi = dict() itov = dict() i = 0 def intern(v): nonlocal i ret = vtoi.get(v) if ret is None: ret = i vtoi[v] = i itov[i] = v i += 1 return ret open_set = {start} open_heapq = [(0, intern(start))] came_from = dict() g_score = defaultdict(lambda: float('inf')) g_score[start] = 0 f_score = {} f_score[start] = h(start) while open_set: while True: f, current = heappop(open_heapq) current = itov[current] if current in open_set: break if callable(goal): if goal(current): return f elif current == goal: return f open_set.remove(current) for neighbour in neighbours(current): tentative_g = g_score[current] + d(current, neighbour) if tentative_g < g_score[neighbour]: came_from[neighbour] = current g_score[neighbour] = tentative_g f_score[neighbour] = tentative_g + h(neighbour) open_set.add(neighbour) heappush(open_heapq, (f_score[neighbour], intern(neighbour))) return None T = TypeVar('T') 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) if len(window) == n: yield tuple(window) for x in it: window.append(x) yield tuple(window) def open_day(n: int): if len(argv) == 2: return open(argv[1]) return open(f'{n}.in') def parse_day(n: int, regex: str): with open_day(n) as f: return re_match(regex, f.read().rstrip()).groups()