80 lines
2.0 KiB
Python
80 lines
2.0 KiB
Python
|
from typing import Iterable
|
||
|
|
|
||
|
|
WORD = "XMAS"
|
||
|
|
def words_from_grid(grid, startx, starty) -> Iterable[str]:
|
||
|
|
num_rows = len(grid)
|
||
|
|
num_cols = len(grid[0])
|
||
|
|
|
||
|
|
for (dx, dy) in (
|
||
|
|
(1, 0), (-1, 0),
|
||
|
|
(0, 1), (0, -1),
|
||
|
|
(1, 1), (-1, 1),
|
||
|
|
(1, -1), (-1, -1)
|
||
|
|
):
|
||
|
|
x = startx
|
||
|
|
y = starty
|
||
|
|
word = ""
|
||
|
|
for i in range(len(WORD)):
|
||
|
|
if x >= num_cols or x < 0:
|
||
|
|
continue
|
||
|
|
if y >= num_rows or y < 0:
|
||
|
|
continue
|
||
|
|
|
||
|
|
word += grid[y][x]
|
||
|
|
x += dx
|
||
|
|
y += dy
|
||
|
|
|
||
|
|
yield word
|
||
|
|
|
||
|
|
def cross_words_from_grid(grid, startx, starty) -> tuple[str, str]:
|
||
|
|
num_rows = len(grid)
|
||
|
|
num_cols = len(grid[0])
|
||
|
|
|
||
|
|
output = []
|
||
|
|
for (dx, dy) in (
|
||
|
|
(1, 1), (1, -1)
|
||
|
|
):
|
||
|
|
x = startx - dx
|
||
|
|
y = starty - dy
|
||
|
|
word = ""
|
||
|
|
for _ in range(3):
|
||
|
|
if x >= num_cols or x < 0:
|
||
|
|
continue
|
||
|
|
if y >= num_rows or y < 0:
|
||
|
|
continue
|
||
|
|
|
||
|
|
word += grid[y][x]
|
||
|
|
x += dx
|
||
|
|
y += dy
|
||
|
|
output.append(word)
|
||
|
|
# o: tuple[str, str] =
|
||
|
|
return tuple(output)
|
||
|
|
|
||
|
|
def task1(grid: list[str]) -> int:
|
||
|
|
count = 0
|
||
|
|
for y in range(len(grid)):
|
||
|
|
for x in range(len(grid[0])):
|
||
|
|
for word in words_from_grid(grid, x, y):
|
||
|
|
if word == WORD or word == str(reversed(WORD)):
|
||
|
|
count += 1
|
||
|
|
return count
|
||
|
|
|
||
|
|
def task2(grid: list[str]) -> int:
|
||
|
|
count = 0
|
||
|
|
for y in range(len(grid)):
|
||
|
|
for x in range(len(grid[0])):
|
||
|
|
leading, anti = cross_words_from_grid(grid, x, y)
|
||
|
|
if (leading == "MAS" or leading == "SAM") and (anti == "MAS" or anti == "SAM"):
|
||
|
|
count += 1
|
||
|
|
return count
|
||
|
|
|
||
|
|
# I can't believe I got both parts first try... this should've been a very error-prone episode!
|
||
|
|
def main():
|
||
|
|
with open("i4.txt") as f:
|
||
|
|
text = f.read().splitlines()
|
||
|
|
print(task1(text))
|
||
|
|
print(task2(text))
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|