speed up 2023/23

aoc_wim/aoc2023/q23.py

@@ -13,8 +13,8 @@
 
 nodes = {}
 dzs = dict(zip("^>v<", [-1j, 1, 1j, -1]))
-for z, g in grid.items():
-    if g == "." and sum(grid.get(z + dz) in dzs for dz in dzs.values()) >= 2:
+for z in grid.z(".", first=False):
+    if sum(grid.get(z + dz) in dzs for dz in dzs.values()) >= 2:
         nodes[z] = {}
 
 
@@ -35,24 +35,29 @@ def find_next_node(z0, dz0):
             z1, dz1, dist = find_next_node(z0, dz0)
             branches[dz0] = z1, dist
             nodes[z1][dz1] = z0, dist
+perimeter = {z for z, branches in nodes.items() if len(branches) < 4}
+
 
 
 def longest_path(part="a"):
     grid[start + 1j] = "v"
-    paths = []
-    stack = [(start, 0)]
+    dists = []
+    stack = [(start, {start}, 0)]
     while stack:
-        path = stack.pop()
-        z0, d0 = path[-2:]
+        z0, visited, d0 = stack.pop()
         if z0 == end:
-            paths.append(path)
+            dists.append(d0)
         branches = nodes[z0]
         for g, dz in dzs.items():
-            if dz in branches and grid[z0 + dz] in ([g] if part == "a" else dzs):
+            if dz in branches:
                 z1, dist = branches[dz]
-                if z1 not in path:
-                    stack.append(path + (z1, d0 + dist))
-    return max(p[-1] for p in paths)
+                if z1 in visited:
+                    continue
+                if grid[z0 + dz] != g:
+                    if part == "a" or {z0, z1} <= perimeter:
+                        continue
+                stack.append((z1, visited.union([z1]), d0 + dist))
+    return max(dists)
 
 
 print("answer_a:", longest_path(part="a"))