advent_of_code_2024/day4/run.py

from typing import List, Tuple

EXAMPLE = """
MMMSXXMASM
MSAMXMSMSA
AMXSXMAAMM
MSAMASMSMX
XMASAMXAMM
XXAMMXXAMA
SMSMSASXSS
SAXAMASAAA
MAMMMXMMMM
MXMXAXMASX
"""


class FindXMAS:
    def __init__(self, input: str) -> None:
        self.input = input
        self.inputgrid = self._to_grid(self.input)

    def __str__(self) -> str:
        return f"Part1: {self._part1}\nPart2: {self._part2}"

    @property
    def _part1(self) -> int:
        return len(self.locate_xmas())

    @property
    def _part2(self) -> int:
        return len(self.locate_x_mas())

    @staticmethod
    def _to_grid(input: str) -> List[List[str]]:
        return [list(line) for line in input.splitlines() if line]

    def find_neighbours(
        self,
        origin: Tuple[int, int],
        target: str = "XMAS",
    ) -> List[Tuple[Tuple[int, ...]]]:
        locations: List[Tuple[Tuple[int, ...]]] = []
        x, y = origin
        if self.inputgrid[x][y].upper() != target[0].upper():
            return []
        for dx in range(-1, 2):
            for dy in range(-1, 2):
                if dx == 0 and dy == 0:
                    continue
                current = [(x, y)]
                for n in range(1, len(target)):
                    nx = x + (n * dx)
                    ny = y + (n * dy)
                    if nx < 0 or ny < 0:
                        continue
                    try:
                        if self.inputgrid[nx][ny].upper() == target[n].upper():
                            current.append((nx, ny))
                    except IndexError:
                        pass
                if len(current) == len(target):
                    locations.append(tuple(current))
        return locations

    def locate_xmas(self) -> List[Tuple[Tuple[int, ...]]]:
        locations: List[Tuple[Tuple[int, ...]]] = []

        for i, line in enumerate(self.inputgrid):
            for j, char in enumerate(line):
                if char.upper() == "X":
                    neighbours = self.find_neighbours((i, j), target="XMAS")
                    if neighbours is None:
                        continue
                    locations = [*locations, *neighbours]

        return locations

    def locate_x_mas(self) -> List[Tuple[Tuple[int, ...]]]:
        mas: List[Tuple[Tuple[int, ...]]] = []

        for i, line in enumerate(self.inputgrid):
            for j, char in enumerate(line):
                if char.upper() == "M":
                    neighbours = [
                        val
                        for val in self.find_neighbours((i, j), target="MAS")
                        if all([len(set(part)) == len("MAS") for part in zip(*val)])
                    ]
                    if neighbours is None:
                        continue
                    mas = [*mas, *neighbours]

        locations: List[Tuple[Tuple[int, ...]]] = []
        # breakpoint()
        for loc in mas:
            oth = [val[1] for val in mas if val != loc and val[::-1] != loc]
            # breakpoint()
            if loc[1] in oth:
                locations.append(loc)

        return mas


if __name__ == "__main__":
    example = FindXMAS(EXAMPLE)
    print(example)
    assert example._part1 == 18
    with open("input.txt", "r") as f:
        puzzle = FindXMAS(f.read())
        print(puzzle)