2024 day 15.

src/holt59/aoc/2024/day15.py

@@ -1,7 +1,213 @@
-from typing import Any, Iterator
+import itertools
+from typing import Any, Callable, Iterator, TypeAlias
+
+import numpy as np
+from numpy.typing import NDArray
 
 from ..base import BaseSolver
 
+Robot: TypeAlias = tuple[int, int]
+Grid: TypeAlias = list[list[str]]
+ImageGrid: TypeAlias = NDArray[np.uint8]
+
 
 class Solver(BaseSolver):
-    def solve(self, input: str) -> Iterator[Any]: ...
+    def print_grid(self, grid: Grid, robot: Robot):
+        if self.verbose:
+            grid = [[c for c in row] for row in grid]
+            grid[robot[0]][robot[1]] = "\033[31;1m@\033[00m"
+            for row in grid:
+                self.logger.info("".join(row))
+
+    def convert_grid(self, grid: Grid, robot: Robot) -> ImageGrid:
+        import numpy as np
+
+        grid = [[c for c in row] for row in grid]
+        grid[robot[0]][robot[1]] = "@"
+        return np.array(
+            [
+                [
+                    # {
+                    #     "#": [255, 255, 255],
+                    #     "@": [255, 0, 0],
+                    #     "O": [0, 255, 0],
+                    #     "[": [0, 0, 255],
+                    #     "]": [0, 0, 255],
+                    #     ".": [0, 0, 0],
+                    # }[col]
+                    ord(col)
+                    for col in row
+                ]
+                for row in grid
+            ],
+            dtype=np.uint8,
+        )
+
+    def step_part1(self, grid: Grid, move: str, robot: Robot):
+        match move:
+            case "^":
+                d_row, d_col = -1, 0
+            case ">":
+                d_row, d_col = 0, 1
+            case "v":
+                d_row, d_col = 1, 0
+            case "<":
+                d_row, d_col = 0, -1
+            case _:
+                assert False
+
+        row, col = robot
+        if grid[row + d_row][col + d_col] == ".":
+            robot = (row + d_row, col + d_col)
+        elif grid[row + d_row][col + d_col] != "#":
+            assert grid[row + d_row][col + d_col] == "O"
+
+            n = 1
+            while grid[row + n * d_row][col + n * d_col] == "O":
+                n += 1
+
+            if grid[row + n * d_row][col + n * d_col] == ".":
+                robot = (row + d_row, col + d_col)
+                grid[row + d_row][col + d_col] = "."
+                for k in range(2, n + 1):
+                    grid[row + k * d_row][col + k * d_col] = "O"
+
+        return grid, robot
+
+    def step_part2(self, grid: Grid, move: str, robot: Robot):
+        match move:
+            case "^":
+                d_row, d_col = -1, 0
+            case ">":
+                d_row, d_col = 0, 1
+            case "v":
+                d_row, d_col = 1, 0
+            case "<":
+                d_row, d_col = 0, -1
+            case _:
+                assert False
+
+        row, col = robot
+        if grid[row + d_row][col + d_col] == ".":
+            robot = (row + d_row, col + d_col)
+        elif grid[row + d_row][col + d_col] == "#":
+            ...
+        elif move in "<>":
+            assert d_row == 0
+            assert grid[row][col + d_col] in "[]"
+
+            n = 1
+            while grid[row][col + n * d_col] in "[]":
+                n += 1
+
+            if grid[row][col + n * d_col] == ".":
+                robot = (row, col + d_col)
+                for k in range(n, 1, -1):
+                    grid[row][col + k * d_col] = grid[row][col + (k - 1) * d_col]
+                grid[row + d_row][col + d_col] = "."
+
+        elif move in "^v":
+            assert d_col == 0
+            assert grid[row + d_row][col + d_col] in "[]"
+
+            n = 1
+            boxes: list[set[int]] = [{col}]
+            while True:
+                to_move = boxes[-1]
+                if any(grid[row + n * d_row][c] == "#" for c in to_move):
+                    break
+                if all(grid[row + n * d_row][c] == "." for c in to_move):
+                    break
+
+                as_move: set[int] = set()
+
+                for c in to_move:
+                    if grid[row + n * d_row][c] == "]":
+                        as_move.update({c - 1, c})
+                    elif grid[row + n * d_row][c] == "[":
+                        as_move.update({c, c + 1})
+
+                boxes.append(as_move)
+                n += 1
+
+            if all(grid[row + n * d_row][c] == "." for c in boxes[-1]):
+                for k, to_move in zip(range(n, 1, -1), boxes[-1:0:-1], strict=True):
+                    for c in to_move:
+                        grid[row + k * d_row][c] = grid[row + (k - 1) * d_row][c]
+                        grid[row + (k - 1) * d_row][c] = "."
+                robot = (row + d_row, col + d_col)
+
+        return grid, robot
+
+    def run(
+        self,
+        grid: Grid,
+        moves: str,
+        fn: Callable[
+            [Grid, str, Robot],
+            tuple[Grid, Robot],
+        ],
+        generate: bool,
+    ) -> tuple[Grid, list[ImageGrid]]:
+        # find robot
+        n_rows, n_cols = len(grid), len(grid[0])
+        robot = next(
+            (i, j) for i in range(n_rows) for j in range(n_cols) if grid[i][j] == "@"
+        )
+        grid[robot[0]][robot[1]] = "."
+
+        # initialize
+        images: list[ImageGrid] = []
+
+        if generate:
+            images.append(self.convert_grid(grid, robot))
+
+        self.print_grid(grid, robot)
+
+        for move in self.progress.wrap(moves):
+            self.logger.debug(f"Move '{move}'...")
+            grid, robot = fn(grid, move, robot)
+
+            if generate:
+                images.append(self.convert_grid(grid, robot))
+
+        self.print_grid(grid, robot)
+
+        return grid, images
+
+    def solve(self, input: str) -> Iterator[Any]:
+        grid_s, moves = input.split("\n\n")
+        moves = "".join(moves.split())
+
+        grid, images = self.run(
+            [[c for c in r] for r in grid_s.splitlines()],
+            moves,
+            self.step_part1,
+            self.files is not None,
+        )
+        yield sum(
+            100 * i_row + i_col
+            for i_row, row in enumerate(grid)
+            for i_col, col in enumerate(row)
+            if col == "O"
+        )
+
+        grid, images = self.run(
+            [
+                list(
+                    itertools.chain.from_iterable(
+                        {"O": "[]", "@": "@.", "#": "##", ".": ".."}[c] for c in r
+                    )
+                )
+                for r in grid_s.splitlines()
+            ],
+            moves,
+            self.step_part2,
+            self.files is not None,
+        )
+        yield sum(
+            100 * i_row + i_col
+            for i_row, row in enumerate(grid)
+            for i_col, col in enumerate(row)
+            if col == "["
+        )