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}")