Poetry stuff.

src/holt59/aoc/2023/day16.py

@@ -0,0 +1,110 @@
+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}")