114 lines
3.3 KiB
Python
114 lines
3.3 KiB
Python
from typing import Any, Iterator, Literal, TypeAlias, cast
|
|
|
|
from ..base import BaseSolver
|
|
|
|
CellType: TypeAlias = Literal[".", "|", "-", "\\", "/"]
|
|
Direction: TypeAlias = Literal["R", "L", "U", "D"]
|
|
|
|
Mappings: dict[
|
|
CellType,
|
|
dict[
|
|
Direction,
|
|
tuple[tuple[tuple[int, int, Direction], ...], tuple[Direction, ...]],
|
|
],
|
|
] = {
|
|
".": {
|
|
"R": (((0, +1, "R"),), ("R", "L")),
|
|
"L": (((0, -1, "L"),), ("R", "L")),
|
|
"U": (((-1, 0, "U"),), ("U", "D")),
|
|
"D": (((+1, 0, "D"),), ("U", "D")),
|
|
},
|
|
"-": {
|
|
"R": (((0, +1, "R"),), ("R", "L")),
|
|
"L": (((0, -1, "L"),), ("R", "L")),
|
|
"U": (((0, +1, "R"), (0, -1, "L")), ("U", "D")),
|
|
"D": (((0, +1, "R"), (0, -1, "L")), ("U", "D")),
|
|
},
|
|
"|": {
|
|
"U": (((-1, 0, "U"),), ("U", "D")),
|
|
"D": (((+1, 0, "D"),), ("U", "D")),
|
|
"R": (((-1, 0, "U"), (+1, 0, "D")), ("R", "L")),
|
|
"L": (((-1, 0, "U"), (+1, 0, "D")), ("R", "L")),
|
|
},
|
|
"/": {
|
|
"R": (((-1, 0, "U"),), ("R", "D")),
|
|
"L": (((+1, 0, "D"),), ("L", "U")),
|
|
"U": (((0, +1, "R"),), ("U", "L")),
|
|
"D": (((0, -1, "L"),), ("R", "D")),
|
|
},
|
|
"\\": {
|
|
"R": (((+1, 0, "D"),), ("R", "U")),
|
|
"L": (((-1, 0, "U"),), ("L", "D")),
|
|
"U": (((0, -1, "L"),), ("U", "R")),
|
|
"D": (((0, +1, "R"),), ("L", "D")),
|
|
},
|
|
}
|
|
|
|
|
|
def propagate(
|
|
layout: list[list[CellType]], start: tuple[int, int], direction: Direction
|
|
) -> list[list[tuple[Direction, ...]]]:
|
|
n_rows, n_cols = len(layout), len(layout[0])
|
|
|
|
beams: list[list[tuple[Direction, ...]]] = [
|
|
[() for _ in range(len(layout[0]))] for _ in range(len(layout))
|
|
]
|
|
|
|
queue: list[tuple[tuple[int, int], Direction]] = [(start, direction)]
|
|
|
|
while queue:
|
|
(row, col), direction = queue.pop()
|
|
|
|
if (
|
|
row not in range(0, n_rows)
|
|
or col not in range(0, n_cols)
|
|
or direction in beams[row][col]
|
|
):
|
|
continue
|
|
|
|
moves, update = Mappings[layout[row][col]][direction]
|
|
|
|
beams[row][col] += update
|
|
|
|
for move in moves:
|
|
queue.append(((row + move[0], col + move[1]), move[2]))
|
|
|
|
return beams
|
|
|
|
|
|
class Solver(BaseSolver):
|
|
def solve(self, input: str) -> Iterator[Any]:
|
|
layout: list[list[CellType]] = [
|
|
[cast(CellType, col) for col in row] for row in input.splitlines()
|
|
]
|
|
|
|
beams = propagate(layout, (0, 0), "R")
|
|
|
|
if self.files:
|
|
self.files.create(
|
|
"beams.txt",
|
|
"\n".join(
|
|
"".join("#" if col else "." for col in row) for row in beams
|
|
).encode(),
|
|
True,
|
|
)
|
|
|
|
# part 1
|
|
yield sum(sum(map(bool, row)) for row in beams)
|
|
|
|
# part 2
|
|
n_rows, n_cols = len(layout), len(layout[0])
|
|
cases: list[tuple[tuple[int, int], Direction]] = []
|
|
|
|
for row in range(n_rows):
|
|
cases.append(((row, 0), "R"))
|
|
cases.append(((row, n_cols - 1), "L"))
|
|
for col in range(n_cols):
|
|
cases.append(((0, col), "D"))
|
|
cases.append(((n_rows - 1, col), "U"))
|
|
|
|
yield max(
|
|
sum(sum(map(bool, row)) for row in propagate(layout, start, direction))
|
|
for start, direction in cases
|
|
)
|