from lib import get_data from lib import Grid2D data = get_data(__file__) g = Grid2D(data) DIRS = { "N": (-1, 0), "E": (0, 1), "S": (1, 0), "W": (0, -1), "NE": (-1, 1), "SE": (1, 1), "SW": (1, -1), "NW": (-1, -1), } t1, t2 = 0, 0 for row in range(g.n_rows): for col in range(g.n_cols): for d in DIRS.values(): coords = [(row + d[0] * i, col + d[1] * i) for i in range(len("XMAS"))] for (r, c), z in zip(coords, "XMAS"): if not (0 <= r < g.n_rows and 0 <= c < g.n_cols and g[(r, c)] == z): break else: t1 += 1 coords = [ (row - 1, col - 1), (row - 1, col + 1), (row, col), (row + 1, col - 1), (row + 1, col + 1), ] for s in ["MMASS", "MSAMS", "SSAMM", "SMASM"]: for (r, c), z in zip(coords, s): if not (0 <= r < g.n_rows and 0 <= c < g.n_cols and g[(r, c)] == z): break else: t2 += 1 print(t1) print(t2)