diff --git a/puzzles/6.py b/puzzles/6.py new file mode 100644 index 0000000..d0a9f97 --- /dev/null +++ b/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)