diff options
-rw-r--r-- | 15.py | 54 | ||||
-rw-r--r-- | 5.py | 6 | ||||
-rw-r--r-- | 9.py | 44 | ||||
-rw-r--r-- | utils.py | 60 |
4 files changed, 109 insertions, 55 deletions
@@ -1,21 +1,34 @@ from collections import defaultdict -from utils import open_day +from utils import open_day, Point2D, adjacent_bounded from heapq import heappop, heappush 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, start)] + open_heapq = [(0, intern(start))] came_from = dict() g_score = defaultdict(lambda: float('inf')) g_score[start] = 0 f_score = defaultdict(lambda: float('inf')) f_score[start] = h(start) - while open_set: - f, current = open_heapq[0] + while True: + f, current = heappop(open_heapq) + current = itov[current] + if current in open_set: break if current == goal: return f - heappop(open_heapq) open_set.remove(current) for neighbour in neighbours(current): tentative_g = g_score[current] + d(current, neighbour) @@ -23,31 +36,22 @@ def a_star(start, goal, neighbours, h, d): came_from[neighbour] = current g_score[neighbour] = tentative_g f_score[neighbour] = tentative_g + h(neighbour) - if neighbour not in open_set: - open_set.add(neighbour) - heappush(open_heapq, (f_score[neighbour], neighbour)) + open_set.add(neighbour) + heappush(open_heapq, (f_score[neighbour], intern(neighbour))) return None with open_day(15) as f: inp = [[int(c) for c in l.rstrip()] for l in f] -def neighbours(p, dim): - for dx, dy in ((-1, 0), (1, 0), (0, -1), (0, 1)): - px, py = p[0] + dx, p[1] + dy - if px < 0 or py < 0 or px >= dim[0] or py >= dim[1]: continue - yield px, py - -def manhattan(a, b): return abs(a[0] - b[0]) + abs(a[1] - b[1]) def d(a, b): - x, y = b - add = x // len(inp[0]) + y // len(inp) - x, y = x % len(inp[0]), y % len(inp) + add = b.x // len(inp[0]) + b.y // len(inp) + x, y = b.x % len(inp[0]), b.y % len(inp) return (inp[y][x] + add - 1) % 9 + 1 -start = (0, 0) -goal1 = (len(inp[0]) - 1, len(inp) - 1) -neighbours1 = lambda p: neighbours(p, (len(inp[0]), len(inp))) -print(a_star(start, goal1, neighbours1, lambda p: manhattan(p, goal1), d)) -goal2 = (len(inp[0]) * 5 - 1, len(inp) * 5 - 1) -neighbours2 = lambda p: neighbours(p, (len(inp[0]) * 5, len(inp) * 5)) -print(a_star(start, goal2, neighbours2, lambda p: manhattan(p, goal2), d)) +start = Point2D(0, 0) +goal1 = Point2D(len(inp[0]) - 1, len(inp) - 1) +neighbours1 = lambda p: adjacent_bounded(p, Point2D(len(inp[0]), len(inp)), False) +print(a_star(start, goal1, neighbours1, lambda p: 0, d)) +goal2 = Point2D(len(inp[0]) * 5 - 1, len(inp) * 5 - 1) +neighbours2 = lambda p: adjacent_bounded(p, Point2D(len(inp[0]) * 5, len(inp) * 5), False) +print(a_star(start, goal2, neighbours2, lambda p: 0, d)) @@ -1,11 +1,9 @@ from collections import defaultdict, namedtuple -from utils import open_day - -Pos = namedtuple('Pos', ['x', 'y']) +from utils import open_day, Point2D vents = [ tuple( - Pos(*map(int, h.split(','))) + Point2D(*map(int, h.split(','))) for h in l.rstrip().split(' -> ') ) for l in open_day(5).read().rstrip().split('\n') @@ -1,45 +1,43 @@ from functools import reduce -from utils import open_day +from utils import open_day, Point2D, adjacent_bounded from collections import deque heatmap = [[int(c) for c in l.rstrip()] for l in open_day(9)] +bounds = Point2D(len(heatmap[0]), len(heatmap)) adj = [(1, 0), (0, 1), (-1, 0), (0, -1)] -def is_low(x, y): - risk = heatmap[y][x] - for dx, dy in adj: - nx = x + dx - ny = y + dy - if nx >= 0 and nx < len(heatmap[0]) and ny >= 0 and ny < len(heatmap): - if heatmap[y+dy][x+dx] <= risk: - return False +def is_low(p): + risk = heatmap[p.y][p.x] + for n in adjacent_bounded(p, bounds, False): + if heatmap[n.y][n.x] <= risk: + return False return True -def basin_size(x, y): +def basin_size(p): seen = set() - queue = deque(((x, y),)) + queue = deque((p,)) size = 0 while queue: - x, y = queue.popleft() - if (x, y) in seen: + p = queue.popleft() + if p in seen: continue - if x < 0 or x >= len(heatmap[0]) or y < 0 or y >= len(heatmap): - continue - seen.add((x, y)) - if heatmap[y][x] == 9: + seen.add(p) + if heatmap[p.y][p.x] == 9: continue size += 1 - for dx, dy in adj: - queue.append((x + dx, y + dy)) + for n in adjacent_bounded(p, bounds, False): + # print(p, n) + queue.append(n) return size totrisk = 0 lows = [] -for y in range(len(heatmap)): - for x in range(len(heatmap[y])): - if is_low(x, y): +for y in range(bounds.y): + for x in range(bounds.x): + p = Point2D(x, y) + if is_low(p): totrisk += 1 + heatmap[y][x] - lows.append(basin_size(x, y)) + lows.append(basin_size(p)) print(totrisk) lows.sort() print(reduce(lambda x, y: x * y, lows[-3:], 1)) @@ -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) |