advent-of-code/src/holt59/aoc/2024/day4.py

import sys

import numpy as np

lines = np.array(list(map(list, sys.stdin.read().strip().splitlines())))
print(lines)
n = lines.shape[0]

answer_1 = 0
for i in range(n):
    line = "".join(lines[i, :])
    answer_1 += line.count("XMAS") + line.count("SAMX")

    column = "".join(lines[:, i])
    answer_1 += column.count("XMAS") + column.count("SAMX")

    diag = "".join(np.diagonal(lines, i))
    answer_1 += diag.count("XMAS") + diag.count("SAMX")

    diag = "".join(np.diagonal(lines[::-1, :], i))
    answer_1 += diag.count("XMAS") + diag.count("SAMX")

    if i != 0:
        diag = "".join(np.diagonal(lines, -i))
        answer_1 += diag.count("XMAS") + diag.count("SAMX")

        diag = "".join(np.diagonal(lines[::-1, :], -i))
        answer_1 += diag.count("XMAS") + diag.count("SAMX")

answer_2 = sum(
    "".join(
        (
            lines[i - 1, j - 1],
            lines[i - 1, j + 1],
            lines[i + 1, j - 1],
            lines[i + 1, j + 1],
        )
    )
    in {"MSMS", "SSMM", "MMSS", "SMSM"}
    for i in range(1, n - 1)
    for j in range(1, n - 1)
    if lines[i, j] == "A"
)

print(f"answer 1 is {answer_1}")
print(f"answer 2 is {answer_2}")