aoc2024day4.py

# === Parsing ===

with open("input.txt", "r") as f:
    grid = [line.strip() for line in f]

print(grid[0])


# === Part 1 ===

DIRECTIONS = [
    (1, 0),
    (1, 1),
    (0, 1),
    (-1, 1),
    (-1, 0),
    (-1, -1),
    (0, -1),
    (1, -1),
]


def find_word(grid, start, direction):
    x_, y_ = start
    dx, dy = direction
    width, height = len(grid[0]), len(grid)
    for delta, letter in enumerate("XMAS"):
        x = x_ + delta * dx
        y = y_ + delta * dy
        if x < 0 or x >= width or y < 0 or y >= height or grid[y][x] != letter:
            return False
    return True


count_xmas = 0
for y in range(len(grid)):
    for x in range(len(grid[0])):
        if grid[y][x] != "X":
            continue
        for direction in DIRECTIONS:
            count_xmas += find_word(grid, (x, y), direction)

print(count_xmas)


# === Part 2 ===

def find_x_mas(grid, anchor):
    """Finds a “MAS” cross with the anchor point as the "A"."""
    x, y = anchor
    return (
        {grid[y+1][x+1], grid[y-1][x-1]} == set("MS")
        and {grid[y+1][x-1], grid[y-1][x+1]} == set("MS")
    )


x_mas_count = 0
for y in range(1, len(grid) - 1):
    for x in range(1, len(grid[0]) - 1):
        if grid[y][x] == "A":
            x_mas_count += find_x_mas(grid, (x, y))

print(x_mas_count)