import os import sys from typing import Literal, TypeAlias, cast VERBOSE = os.getenv("AOC_VERBOSE") == "True" 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 = [(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 layout: list[list[CellType]] = [ [cast(CellType, col) for col in row] for row in sys.stdin.read().splitlines() ] beams = propagate(layout, (0, 0), "R") if VERBOSE: print("\n".join(["".join("#" if col else "." for col in row) for row in beams])) # part 1 answer_1 = sum(sum(map(bool, row)) for row in beams) print(f"answer 1 is {answer_1}") # 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")) answer_2 = max( sum(sum(map(bool, row)) for row in propagate(layout, start, direction)) for start, direction in cases ) print(f"answer 2 is {answer_2}")