Make day 6 a little nicer

2024/d6.py

@@ -4,52 +4,48 @@ from lib import Grid2D
 data = get_data(__file__)
 
 
-DIRS = {
-    "N": (-1, 0),
-    "E": (0, 1),
-    "S": (1, 0),
-    "W": (0, -1),
-}
-DIRSL = "NESW"
+DIRS = [
+    (-1, 0),
+    (0, 1),
+    (1, 0),
+    (0, -1),
+]
 
 g = Grid2D(data)
 
 t = 0
 for rr in range(g.n_rows):
     for cc in range(g.n_cols):
-        (p,) = g.find("^")
-
+        (p,), dir = g.find("^"), 0
         seen = set()
-        dirl = "N"
-        dir = DIRS["N"]
-
         loop = False
 
-        if g[(rr, cc)] == "#" or g[(rr, cc)] == "^":
+        if g[(rr, cc)] == "#":
             continue
 
-        g[(rr, cc)] = "#"
+        if g[(rr, cc)] != "^":
+            g[(rr, cc)] = "#"
 
         while True:
             if (p, dir) in seen:
                 loop = True
                 break
             seen.add((p, dir))
-            r, c = p[0] + dir[0], p[1] + dir[1]
+            r, c = p[0] + DIRS[dir][0], p[1] + DIRS[dir][1]
             if not (0 <= r < g.n_rows and 0 <= c < g.n_cols):
                 break
 
             if g[(r, c)] == "#":
-                dirl = DIRSL[(DIRSL.index(dirl) + 1) % 4]
-                dir = DIRS[dirl]
-
+                dir = (dir + 1) % 4
             else:
                 p = (r, c)
 
-        g[(rr, cc)] = "."
+        if g[(rr, cc)] == "^":
+            print(len(set(pos for pos, _ in seen)))
+        else:
+            g[(rr, cc)] = "."
 
         if loop:
             t += 1
 
-print(len(set(pos for pos, _ in seen)))
 print(t)