advent-of-code-2024/day_04/part_a.py

79 lines
2.1 KiB
Python

import aocd
from enum import Enum
import collections
class Direction(Enum):
NORTH = (0, -1)
NORTHEAST = (1, -1)
EAST = (1, 0)
SOUTHEAST = (1, 1)
SOUTH = (0, 1)
SOUTHWEST = (-1, 1)
WEST = (-1, 0)
NORTHWEST = (-1, -1)
class Letter:
target = 'XMAS'
def __init__(self, char, x, y, grid):
self.char = char
self.x = x
self.y = y
self.grid = grid
def neighbour_coords(self, direction):
return (self.x + direction.value[0], self.y + direction.value[1])
def neighbour(self, direction):
return self.grid[self.neighbour_coords(direction)]
def num_starts(self):
"""Return number of times this letter is the start of the target word"""
return sum(self._is_nth_position(0, d) for d in Direction)
def _is_nth_position(self, n, direction):
"""
Return number of times this letter is the nth letter of the target word
in the given direction
"""
# Must be the correct letter
if self.char != self.target[n]:
return False
# Return True if it's the last letter
if n == len(self.target) - 1:
return True
# Return False if there are no more letter in this direction
if (next_letter := self.neighbour(direction)) is None:
return False
# Ask the next letter
return next_letter._is_nth_position(n+1, direction)
def parse(data):
"""
Parse string into grid, return all possible first letters of XMAS
"""
grid = collections.defaultdict(lambda: None)
x_list = []
for y, r in enumerate(data.split('\n')):
for x, c in enumerate(r):
grid[(x, y)] = Letter(c, x, y, grid)
if c == 'X':
x_list.append(grid[(x, y)])
return x_list
def solve(data):
x_list = parse(data)
return sum(x.num_starts() for x in x_list)
if __name__ == '__main__':
solution = solve(aocd.get_data(year=2024, day=4))
print(solution)
aocd.submit(solution, year=2024, day=4, part='a')