AoC 2024 Day 4 -- Part 1

from pathlib import Path


def part_one(input: str) -> int:
    xmas = "XMAS"

    if Path(input).exists():
        with open(input) as f:
            input = f.read()

    check_position = {}
    count_words = 0
    directions = [(0, 1), (1, -1), (1, 0), (1, 1)]

    for i, line in enumerate(input.strip().splitlines()):
        for j, char in enumerate(line):
            to_be_checked = check_position.pop((i, j), [])

            for schar, coord, dir in to_be_checked:
                if schar == char:
                    if char in "XS":
                        count_words += 1
                    else:
                        next_pos = (i + coord[0], j + coord[1])
                        next_char = xmas[xmas.find(char) + dir]
                        check_position.setdefault(next_pos, []).append(
                            (next_char, coord, dir)
                        )

            if char in "XS":
                direction_shift = 1 if char == "X" else -1
                for dx, dy in directions:
                    next_pos = (i + dx, j + dy)
                    next_char = xmas[xmas.find(char) + direction_shift]
                    check_position.setdefault(next_pos, []).append(
                        (next_char, (dx, dy), direction_shift)
                    )

        check_position = {
            key: val for key, val in check_position.items() if key[0] > i
        }

    return count_words