2024-04 python cleanup

1 files changed, 22 insertions, 42 deletions

diff --git a/2024/04/python/main.py b/2024/04/python/main.py
index a16a943..cc04f04 100644
--- a/2024/04/python/main.py
+++ b/2024/04/python/main.py
@@ -1,47 +1,27 @@
-import sys
-
-with open(sys.argv[1], "r") as f:
-    grid = [line.strip() for line in f]
-
-min_x = 0
-min_y = 0
-max_x = len(grid)
-max_y = len(grid[0])
-
-
-def count_xmas(sx, sy):
-    total = 0
-    for dx in [-1, 0, 1]:
-        for dy in [-1, 0, 1]:
-            for n, char in enumerate("XMAS"):
-                x = sx + n * dx
-                y = sy + n * dy
-                if not (min_x <= x < max_x and min_y <= y < max_y):
-                    break
-                if grid[x][y] != char:
-                    break
-            else:
-                total += 1
-    return total
-
-
-def count_x_mas(sx, sy):
-    if sx == min_x or sx == max_x - 1 or sy == min_y or sy == max_y - 1:
-        return 0
-    if grid[sx][sy] != "A":
-        return 0
-    mmss = "".join(
-        (
-            grid[sx - 1][sy - 1],
-            grid[sx - 1][sy + 1],
-            grid[sx + 1][sy + 1],
-            grid[sx + 1][sy - 1],
-        ),
+from fileinput import input
+
+grid = {
+    complex(idx, idy): c
+    for idy, line in enumerate(input())
+    for idx, c in enumerate(line.strip())
+}
+dirs = {complex(dx, dy) for dx in [-1, 0, 1] for dy in [-1, 0, 1]}
+
+
+def xmas(x, v):
+    return all(grid.get(x + n * v) == c for n, c in enumerate("XMAS"))
+
+
+def x_mas(x, v):
+    return (
+        grid.get(x) == "A"
+        and grid.get(x + v) == grid.get(x + 1j * v) == "M"
+        and grid.get(x - v) == grid.get(x - 1j * v) == "S"
     )
-    return mmss in {"MMSS", "MSSM", "SSMM", "SMMS"}
 
 
-silver = sum(count_xmas(x, y) for x in range(max_x) for y in range(max_y))
-gold = sum(count_x_mas(x, y) for x in range(max_x) for y in range(max_y))
+silver = sum(xmas(x, v) for x in grid.keys() for v in dirs)
+gold = sum(x_mas(x, v) for x in grid.keys() for v in dirs if abs(v) > 1)
+
 print("silver:", silver)
 print("gold:", gold)