Day 16 (#23)

* Solve 16.a

* Solve 16.b

aoc_2024/day_16/a.py

@@ -0,0 +1,90 @@
+from dataclasses import dataclass
+from functools import cached_property
+from queue import PriorityQueue
+from typing import Generator
+from aoc_2024.day_16.parser import Parser
+from aoc_2024.utils.point import Point
+
+
+DIRECTIONS = [
+    Point(1, 0),
+    Point(0, 1),
+    Point(-1, 0),
+    Point(0, -1),
+]
+
+MOVE_COST = 1
+ROTATE_COST = 1000
+
+
+@dataclass
+class Day16PartASolver:
+    maze: list[list[str]]
+
+    @property
+    def solution(self) -> int:
+        visited = set[tuple[Point, int]]()
+        queue = PriorityQueue[tuple[int, Point, int]]()
+        queue.put((0, self.start, 0))
+
+        while not queue.empty():
+            score, pos, direction_index = queue.get()
+            if (pos, direction_index) in visited:
+                continue
+
+            if pos == self.end:
+                return score
+
+            visited.add((pos, direction_index))
+            for next_state in self.get_next_states(
+                score=score,
+                pos=pos,
+                direction_index=direction_index,
+            ):
+                _, next_pos, next_direction_index = next_state
+                if (next_pos, next_direction_index) not in visited:
+                    queue.put(next_state)
+        assert False, "uh oh"
+
+    def get_next_states(
+        self,
+        score: int,
+        pos: Point,
+        direction_index: int,
+    ) -> Generator[tuple[int, Point, int]]:
+        move_forward = pos.add(DIRECTIONS[direction_index])
+        if self.maze[move_forward.y][move_forward.x] != "#":
+            yield score + MOVE_COST, move_forward, direction_index
+        yield score + ROTATE_COST, pos, (direction_index - 1) % 4
+        yield score + ROTATE_COST, pos, (direction_index + 1) % 4
+
+    @cached_property
+    def start(self) -> Point:
+        x = 1
+        y = len(self.maze) - 2
+        assert self.maze[y][x] == "S"
+        return Point(x, y)
+
+    @cached_property
+    def end(self) -> Point:
+        x = len(self.maze[0]) - 2
+        y = 1
+        assert self.maze[y][x] == "E"
+        return Point(x, y)
+
+
+def solve(input: str) -> int:
+    data = Parser.parse(input)
+    solver = Day16PartASolver(data)
+
+    return solver.solution
+
+
+def get_solution() -> int:
+    with open("aoc_2024/day_16/input.txt", "r") as f:
+        input = f.read()
+    return solve(input)
+
+
+if __name__ == "__main__":
+    print(get_solution())

aoc_2024/day_16/b.py

@@ -0,0 +1,107 @@
+from dataclasses import dataclass
+from functools import cached_property
+from queue import PriorityQueue
+from typing import Generator
+from aoc_2024.day_16.parser import Parser
+from aoc_2024.utils.point import Point
+
+
+DIRECTIONS = [
+    Point(1, 0),
+    Point(0, 1),
+    Point(-1, 0),
+    Point(0, -1),
+]
+
+MOVE_COST = 1
+ROTATE_COST = 1000
+
+
+@dataclass
+class Day16PartBSolver:
+    maze: list[list[str]]
+
+    @property
+    def solution(self) -> int:
+        queue = PriorityQueue[tuple[int, Point, int, tuple[Point, ...]]]()
+        queue.put((0, self.start, 0, tuple()))
+        best_score: None | int = None
+        in_best_path = {self.start, self.end}
+        best_score_to_state: dict[tuple[Point, int], int] = {}
+
+        while not queue.empty():
+            score, pos, direction_index, path = queue.get()
+            if best_score is not None and score > best_score:
+                break
+
+            key = (pos, direction_index)
+            if key in best_score_to_state:
+                if score > best_score_to_state[key]:
+                    continue
+            else:
+                best_score_to_state[key] = score
+
+            if pos == self.end:
+                best_score = score
+                for p in path:
+                    in_best_path.add(p)
+
+            for next_state in self.get_next_states(
+                score=score,
+                pos=pos,
+                direction_index=direction_index,
+            ):
+                next_score, next_pos, next_direction_index = next_state
+                new_key = (next_pos, next_direction_index)
+                if (
+                    new_key not in best_score_to_state
+                    or next_score <= best_score_to_state[new_key]
+                ):
+                    queue.put(
+                        (next_score, next_pos, next_direction_index, path + (pos,))
+                    )
+
+        return len(in_best_path)
+
+    def get_next_states(
+        self,
+        score: int,
+        pos: Point,
+        direction_index: int,
+    ) -> Generator[tuple[int, Point, int]]:
+        move_forward = pos.add(DIRECTIONS[direction_index])
+        if self.maze[move_forward.y][move_forward.x] != "#":
+            yield score + MOVE_COST, move_forward, direction_index
+        yield score + ROTATE_COST, pos, (direction_index - 1) % 4
+        yield score + ROTATE_COST, pos, (direction_index + 1) % 4
+
+    @cached_property
+    def start(self) -> Point:
+        x = 1
+        y = len(self.maze) - 2
+        assert self.maze[y][x] == "S"
+        return Point(x, y)
+
+    @cached_property
+    def end(self) -> Point:
+        x = len(self.maze[0]) - 2
+        y = 1
+        assert self.maze[y][x] == "E"
+        return Point(x, y)
+
+
+def solve(input: str) -> int:
+    data = Parser.parse(input)
+    solver = Day16PartBSolver(data)
+
+    return solver.solution
+
+
+def get_solution() -> int:
+    with open("aoc_2024/day_16/input.txt", "r") as f:
+        input = f.read()
+    return solve(input)
+
+
+if __name__ == "__main__":
+    print(get_solution())

aoc_2024/day_16/parser.py

@@ -0,0 +1,5 @@
+class Parser:
+    @staticmethod
+    def parse(input: str) -> list[list[str]]:
+        lines = input.strip().splitlines()
+        return [list(line) for line in lines]

aoc_2024/utils/point.py

@@ -1,13 +1,21 @@
 from __future__ import annotations
 from dataclasses import dataclass
+from functools import total_ordering
+from typing import Any
 
 
+@total_ordering
 @dataclass(frozen=True)
 class Point:
     x: int
     y: int
 
-    @property
+    def __lt__(self, other: Any) -> bool:
+        if not isinstance(other, Point):
+            return False
+        else:
+            return abs(self.x) + abs(self.y) < abs(other.x) + abs(other.y)
+
     def unit(self) -> Point:
         x = 0 if self.x == 0 else self.x // abs(self.x)
         y = 0 if self.y == 0 else self.y // abs(self.y)

tests/test_day_16/test_a.py

@@ -0,0 +1,24 @@
+import pytest
+from aoc_2024.day_16.a import get_solution, solve
+from aoc_2024.day_16.from_prompt import (
+    SAMPLE_DATA_1,
+    SAMPLE_DATA_2,
+    SAMPLE_SOLUTION_A1,
+    SAMPLE_SOLUTION_A2,
+    SOLUTION_A,
+)
+
+
+@pytest.mark.parametrize(
+    ("input", "expected"),
+    [
+        (SAMPLE_DATA_1, SAMPLE_SOLUTION_A1),
+        (SAMPLE_DATA_2, SAMPLE_SOLUTION_A2),
+    ],
+)
+def test_solve(input: str, expected: int):
+    assert solve(input) == expected
+
+
+def test_my_solution():
+    assert get_solution() == SOLUTION_A

tests/test_day_16/test_b.py

@@ -0,0 +1,24 @@
+import pytest
+from aoc_2024.day_16.b import get_solution, solve
+from aoc_2024.day_16.from_prompt import (
+    SAMPLE_DATA_1,
+    SAMPLE_DATA_2,
+    SAMPLE_SOLUTION_B1,
+    SAMPLE_SOLUTION_B2,
+    SOLUTION_B,
+)
+
+
+@pytest.mark.parametrize(
+    ("input", "expected"),
+    [
+        (SAMPLE_DATA_1, SAMPLE_SOLUTION_B1),
+        (SAMPLE_DATA_2, SAMPLE_SOLUTION_B2),
+    ],
+)
+def test_solve(input: str, expected: int):
+    assert solve(input) == expected
+
+
+def test_my_solution():
+    assert get_solution() == SOLUTION_B