advent-of-code-2024/day_06/part_b.py

75 lines
2.0 KiB
Python

import aocd
import collections
from enum import Enum, nonmember
class Direction(Enum):
NORTH = (0, -1)
WEST = (-1, 0)
SOUTH = (0, 1)
EAST = (1, 0)
@nonmember
def clockwise_from(d):
return list(Direction)[list(Direction).index(d) - 1]
def facing(coords, direction):
return tuple(map(sum, zip(coords, direction.value)))
def parse(data):
"""
Given maze data, return
(Tuple of guard coords, set of obstacle coords, grid width, grid height)
"""
guard = None
obstacles = set()
for y, row in enumerate(data.split('\n')):
for x, char in enumerate(row):
if char == '#':
obstacles.add((x, y))
elif char == '^':
guard = (x, y)
return (guard, obstacles, x, y)
def simulate(guard, obstacles, width, height, return_visited_tiles=False):
guard_direction = Direction.NORTH
visited = collections.defaultdict(set)
while (0 <= guard[0] <= width) and (0 <= guard[1] <= height):
visited[guard].add(guard_direction)
while (facing_tile := facing(guard, guard_direction)) in obstacles:
guard_direction = Direction.clockwise_from(guard_direction)
# Loop detection, can only be true in modified grids
if guard_direction in visited[facing_tile]:
return True
guard = facing_tile
# No loop found
if return_visited_tiles:
return set(k for k,v in visited.items() if v)
return False
def solve(data):
guard, obstacles, width, height = parse(data)
# Only bother checking positions that the guard crosses
potential_positions = simulate(guard, obstacles, width, height, return_visited_tiles=True)
return sum(simulate(guard, obstacles | {p}, width, height) for p in potential_positions)
if __name__ == '__main__':
puz = aocd.models.Puzzle(year=2024, day=6)
# data = puz.examples[0].input_data
data = puz.input_data
solution = solve(data)
print(solution)
aocd.submit(solution, year=2024, day=6, part='b')