Implement day 6 solver

puzzles/6.py

@@ -0,0 +1,123 @@
+from typing import NamedTuple
+
+
+test_input = """
+....#.....
+.........#
+..........
+..#.......
+.......#..
+..........
+.#..^.....
+........#.
+#.........
+......#...
+""".strip()
+
+test_solution_p1 = 41
+test_solution_p2 = 6
+
+
+def solve_p1(puzzle_input: str) -> int:
+    map_info = _parse_map(puzzle_input)
+    traversal_analysis = _traverse_map(map_info)
+    unique_locations = set((a, b) for a, b, _ in traversal_analysis.visited)
+    return len(unique_locations)
+
+
+def solve_p2(puzzle_input: str) -> int:
+    map_info = _parse_map(puzzle_input)
+    added_obstacles_candidates = set()
+
+    for i, line in enumerate(map_info.map):
+        for j, cell in enumerate(line):
+            if not cell:
+                map_info.map[i][j] = True
+                if _traverse_map(map_info).stuck:
+                    added_obstacles_candidates.add((i, j))
+                map_info.map[i][j] = False
+
+    # Remove guard location from candidates
+    added_obstacles_candidates.discard((map_info.guard[:2]))
+
+    return len(added_obstacles_candidates)
+
+
+class MapInfo(NamedTuple):
+    map: list[list[bool]]
+    guard: tuple[int, int, str]
+
+
+def _parse_map(puzzle_input: str) -> MapInfo:
+    puzzle_map = []
+    guard = None
+
+    for x, line in enumerate(puzzle_input.split("\n")):
+        if not line:
+            continue
+        line_list = list()
+        for y, cell in enumerate(line):
+            if cell in "^v<>":
+                guard = (x, y, cell)
+            line_list.append(cell == "#")
+        puzzle_map.append(line_list)
+
+    assert guard is not None, "Guard not found in map"
+    assert all(
+        len(line) == len(puzzle_map[0]) for line in puzzle_map
+    ), "Map is not rectangular"
+
+    return MapInfo(puzzle_map, guard)
+
+
+class TraversalAnalysis(NamedTuple):
+    visited: set[tuple[int, int]]
+    stuck: bool
+
+
+_rotation_map = {
+    "^": ">",
+    "v": "<",
+    "<": "^",
+    ">": "v",
+}
+
+
+def _traverse_map(map_info: MapInfo) -> TraversalAnalysis:
+    x, y, direction = map_info.guard
+    visited = set()
+    visited.add((x, y, direction))
+
+    while True:
+        # Determine next cell
+        match direction:
+            case "^":
+                next_x, next_y = x - 1, y
+            case "v":
+                next_x, next_y = x + 1, y
+            case "<":
+                next_x, next_y = x, y - 1
+            case ">":
+                next_x, next_y = x, y + 1
+
+        # Guard patrols out of bounds
+        if (
+            next_x < 0
+            or next_x >= len(map_info.map)
+            or next_y < 0
+            or next_y >= len(map_info.map[0])
+        ):
+            break
+
+        # Turn right if an obstacle is encountered
+        if map_info.map[next_x][next_y]:
+            direction = _rotation_map[direction]
+
+        # Otherwise, move to next cell
+        else:
+            x, y = next_x, next_y
+            if (x, y, direction) in visited:
+                return TraversalAnalysis(visited, True)
+            visited.add((x, y, direction))
+
+    return TraversalAnalysis(visited, False)