2022: d24: ex2: add solution

2022/d24/ex2/ex2.py

@@ -0,0 +1,164 @@
+#!/usr/bin/env python
+
+import dataclasses
+import enum
+import sys
+from collections import defaultdict, deque
+from collections.abc import Iterator
+from typing import NamedTuple
+
+
+class Point(NamedTuple):
+    x: int
+    y: int
+
+    def __add__(self, other):
+        if not isinstance(other, Point):
+            return NotImplemented
+        return Point(self.x + other.x, self.y + other.y)
+
+    def __sub__(self, other):
+        if not isinstance(other, Point):
+            return NotImplemented
+        return Point(self.x - other.x, self.y - other.y)
+
+
+class Direction(str, enum.Enum):
+    UP = "^"
+    DOWN = "v"
+    LEFT = "<"
+    RIGHT = ">"
+
+    def to_delta(self) -> Point:
+        match self:
+            case Direction.UP:
+                return Point(-1, 0)
+            case Direction.DOWN:
+                return Point(1, 0)
+            case Direction.LEFT:
+                return Point(0, -1)
+            case Direction.RIGHT:
+                return Point(0, 1)
+
+
+@dataclasses.dataclass
+class ValleyMap:
+    start: Point
+    goal: Point
+    valley_corners: tuple[Point, Point]
+    tornadoes: dict[Point, Direction]
+
+    @classmethod
+    def from_input(cls, input: list[str]) -> "ValleyMap":
+        tornadoes: dict[Point, Direction] = {}
+        for x, line in enumerate(input, start=1):
+            for y, c in enumerate(line, start=1):
+                if c in ("#", "."):
+                    continue
+                tornadoes[Point(x, y)] = Direction(c)
+        return cls(
+            # Start position is always above the upper left corner of valley
+            start=Point(1, 2),
+            # Goal position is always under the lower left corner of valley
+            goal=Point(len(input), len(input[0]) - 1),
+            # Valley is surrounded by walls, except entrance and exit
+            valley_corners=(Point(2, 2), Point(len(input) - 1, len(input[0]) - 1)),
+            tornadoes=tornadoes,
+        )
+
+    def _is_in_valley(self, p: Point) -> bool:
+        # Valley also includes start/end
+        if p in (self.start, self.goal):
+            return True
+        # Otherwise, just do a bounds check for inside the walls
+        ((minx, miny), (maxx, maxy)) = self.valley_corners
+        return (minx <= p.x <= maxx) and (miny <= p.y <= maxy)
+
+    def _wrap_tornado(self, p: Point) -> Point:
+        if self._is_in_valley(p):
+            return p
+        x, y = p
+        h = self.valley_corners[1].x - self.valley_corners[0].x + 1
+        w = self.valley_corners[1].y - self.valley_corners[0].y + 1
+        if x == 1:
+            x += h
+        if y == 1:
+            y += w
+        if x > self.valley_corners[1].x:
+            x -= h
+        if y > self.valley_corners[1].y:
+            y -= w
+        return Point(x, y)
+
+    def navigate(self) -> int:
+        TornadoesMap = dict[Point, list[Direction]]
+
+        def move_tornadoes(map: TornadoesMap) -> dict[Point, list[Direction]]:
+            res: dict[Point, list[Direction]] = defaultdict(list)
+            for p, tornadoes in map.items():
+                for t in tornadoes:
+                    new_pos = self._wrap_tornado(p + t.to_delta())
+                    res[new_pos].append(t)
+            return dict(res)
+
+        def moves(p: Point) -> Iterator[Point]:
+            yield p
+            for dx, dy in ((-1, 0), (1, 0), (0, -1), (0, 1)):
+                yield p + Point(dx, dy)
+
+        def bfs(
+            start: Point, goal: Point, tornadoes: TornadoesMap
+        ) -> tuple[int, TornadoesMap]:
+            # Do a BFS to find the fastest route
+            queue: deque[tuple[int, Point]] = deque([(0, start)])
+            seen: set[tuple[int, Point]] = set()
+            tornado_history = [tornadoes]
+            while queue:
+                dist, pos = queue.popleft()
+                # If goal found, return total distance
+                if pos == goal:
+                    return dist, tornado_history[dist]
+                # Check that we don't do redundant work
+                if (dist, pos) in seen:
+                    continue
+                seen.add((dist, pos))
+                if len(tornado_history) <= (dist + 1):
+                    tornado_history.append(move_tornadoes(tornado_history[-1]))
+                for new_pos in moves(pos):
+                    # Can't move into the walls, but can move in start/end
+                    if not self._is_in_valley(new_pos):
+                        continue
+                    # Can't occupy same space as tornadoes
+                    if new_pos in tornado_history[dist + 1]:
+                        continue
+                    # Enqueue this move to the search space
+                    queue.append((dist + 1, new_pos))
+            assert False  # Sanity check
+
+        tornadoes = {p: [t] for p, t in self.tornadoes.items()}
+        total = 0
+        for start, end in (
+            # First travel
+            (self.start, self.goal),
+            # Back for snacks
+            (self.goal, self.start),
+            # Second travel
+            (self.start, self.goal),
+        ):
+            dist, tornadoes = bfs(start, end, tornadoes)
+            total += dist
+        return total
+
+
+def solve(input: list[str]) -> int:
+    valley = ValleyMap.from_input(input)
+    return valley.navigate()
+
+
+def main() -> None:
+    input = sys.stdin.read().splitlines()
+    print(solve(input))
+
+
+if __name__ == "__main__":
+    main()