161 lines
4.5 KiB
Python
161 lines
4.5 KiB
Python
from dataclasses import dataclass
|
|
from textwrap import dedent
|
|
from typing import Literal, NamedTuple, get_args, override
|
|
from unittest import TestCase
|
|
|
|
from puzzles._solver import Solver
|
|
|
|
direction_map = {
|
|
"^": (-1, 0),
|
|
"v": (1, 0),
|
|
"<": (0, -1),
|
|
">": (0, 1),
|
|
}
|
|
|
|
|
|
class GuardPosition(NamedTuple):
|
|
x: int
|
|
y: int
|
|
direction: tuple[int, int]
|
|
|
|
|
|
def move_position(position: GuardPosition) -> GuardPosition:
|
|
return GuardPosition(
|
|
position.x + position.direction[0],
|
|
position.y + position.direction[1],
|
|
position.direction,
|
|
)
|
|
|
|
|
|
def rotate_position(position: GuardPosition) -> GuardPosition:
|
|
return GuardPosition(
|
|
position.x,
|
|
position.y,
|
|
(position.direction[1], -position.direction[0]),
|
|
)
|
|
|
|
|
|
class TestGuardPosition(TestCase):
|
|
def test(self):
|
|
guard = GuardPosition(0, 0, (0, 1))
|
|
self.assertEqual(move_position(guard), GuardPosition(0, 1, (0, 1)))
|
|
self.assertEqual(rotate_position(guard), GuardPosition(0, 0, (1, 0)))
|
|
|
|
|
|
class TraversalAnalysis(NamedTuple):
|
|
visited: set[GuardPosition]
|
|
stuck: bool
|
|
|
|
|
|
def unique_locations(traversal_analysis: TraversalAnalysis) -> set[tuple[int, int]]:
|
|
return set((p.x, p.y) for p in traversal_analysis.visited)
|
|
|
|
|
|
class TestTraversalAnalysis(TestCase):
|
|
def test(self):
|
|
analysis = TraversalAnalysis(
|
|
{GuardPosition(0, 0, (0, 1)), GuardPosition(2, 5, (1, 0))},
|
|
False,
|
|
)
|
|
self.assertEqual(
|
|
unique_locations(analysis),
|
|
{(0, 0), (2, 5)},
|
|
)
|
|
|
|
|
|
class DaySixSolver(Solver):
|
|
puzzle_map: list[list[bool]]
|
|
guard_origin: GuardPosition
|
|
|
|
@override
|
|
def __init__(self, puzzle_input: str):
|
|
self.puzzle_map = []
|
|
guard_origin = None
|
|
|
|
for x, line in enumerate(puzzle_input.strip().split("\n")):
|
|
line_list = list()
|
|
for y, cell in enumerate(line):
|
|
if cell in direction_map:
|
|
guard_origin = GuardPosition(x, y, direction_map[cell])
|
|
line_list.append(cell == "#")
|
|
self.puzzle_map.append(line_list)
|
|
|
|
if guard_origin is None:
|
|
raise ValueError("Guard not found in map")
|
|
|
|
self.guard_origin = guard_origin
|
|
|
|
if not all(len(line) == len(self.puzzle_map[0]) for line in self.puzzle_map):
|
|
raise ValueError("Map is not rectangular")
|
|
|
|
@override
|
|
def solve_p1(self) -> int:
|
|
traversal_analysis = self.traverse_map()
|
|
return len(unique_locations(traversal_analysis))
|
|
|
|
@override
|
|
def solve_p2(self) -> int:
|
|
candidate_locations = unique_locations(self.traverse_map())
|
|
candidate_locations.discard((self.guard_origin.x, self.guard_origin.y))
|
|
verified_locations = set()
|
|
|
|
for x, y in candidate_locations:
|
|
if not self.puzzle_map[x][y]:
|
|
self.puzzle_map[x][y] = True
|
|
if self.traverse_map().stuck:
|
|
verified_locations.add((x, y))
|
|
self.puzzle_map[x][y] = False
|
|
|
|
return len(verified_locations)
|
|
|
|
def out_of_bounds(self, guard_position: GuardPosition) -> bool:
|
|
return (
|
|
guard_position.x < 0
|
|
or guard_position.x >= len(self.puzzle_map)
|
|
or guard_position.y < 0
|
|
or guard_position.y >= len(self.puzzle_map[0])
|
|
)
|
|
|
|
def traverse_map(self) -> TraversalAnalysis:
|
|
position = self.guard_origin
|
|
visited = set({position})
|
|
|
|
while True:
|
|
next_position = move_position(position)
|
|
if self.out_of_bounds(next_position):
|
|
return TraversalAnalysis(visited, False)
|
|
|
|
# Obstacle is encountered
|
|
if self.puzzle_map[next_position.x][next_position.y]:
|
|
next_position = rotate_position(position)
|
|
|
|
# Guard is stuck
|
|
if next_position in visited:
|
|
return TraversalAnalysis(visited, True)
|
|
|
|
position = next_position
|
|
visited.add(position)
|
|
|
|
|
|
class TestDaySixSolver(TestCase):
|
|
def test(self):
|
|
solver = DaySixSolver(
|
|
dedent(
|
|
"""
|
|
....#.....
|
|
.........#
|
|
..........
|
|
..#.......
|
|
.......#..
|
|
..........
|
|
.#..^.....
|
|
........#.
|
|
#.........
|
|
......#...
|
|
"""
|
|
)
|
|
)
|
|
self.assertEqual(solver.guard_origin, GuardPosition(6, 4, (-1, 0)))
|
|
self.assertEqual(solver.solve_p1(), 41)
|
|
self.assertEqual(solver.solve_p2(), 6)
|