day 12 python tweaks

day12/sol.py

@@ -1,52 +1,59 @@
-data = open("input").read().splitlines(False)
-
 import heapq
 import string
+
 chars = 'ES'+string.ascii_lowercase
 
+data = open("input").read().splitlines(False)
 data = [list(map(chars.index, line)) for line in data]
 #print(data)
+
 def get(x, y):
     return data[y][x]
 
-q = []
+extra = []
 for x in range(len(data[0])):
     for y in range(len(data)):
-        if get(x,y) == 1:
+        if get(x,y) == 1: # 'S'
             data[y][x] = chars.index('a')
             start = (x,y)
-        elif get(x,y) == 0:
+        elif get(x,y) == 0: # 'E'
             data[y][x] = chars.index('z')
             goal = (x,y)
-        elif get(x,y) == 2:
-            # part 2
-            q.append((0,(x,y)))
-            pass
+        elif get(x,y) == 2: # 'a'
+            # for part 2
+            extra.append((x,y))
 
-q.append((0, start))
-path = {}
-while q:
-    cost, (x,y) = heapq.heappop(q)
-    if (x,y) in path:
-        continue
-    #print(x,y,cost)
-    if (x,y) == goal:
-        print('found', cost)
-        break
-    here = get(x,y)
-    path[(x,y)] = cost
-    
-    n = []
-    def check(x,y):
-        if 0 <= y < len(data):
-            if 0 <= x < len(data[y]):
-                if get(x,y) <= here+1 and (x,y) not in path:
-                    n.append((cost+1, (x,y)))
-    check(x-1, y)
-    check(x+1,y)
-    check(x, y-1)
-    check(x, y+1)
+def solve(part):
+    q = [(0, start)]
 
-    for p in n:
-        heapq.heappush(q, p)
+    if part == 2:
+        q.extend((0,p) for p in extra)
 
+    done = {}
+    while q:
+        cost, (x,y) = heapq.heappop(q)
+        if (x,y) in done:
+            continue
+        #print(x,y,cost)
+        if (x,y) == goal:
+            print('found', cost)
+            break
+        here = get(x,y)
+        done[(x,y)] = cost
+
+        n = []
+        def check(x,y):
+            if 0 <= y < len(data):
+                if 0 <= x < len(data[y]):
+                    if get(x,y) <= here+1 and (x,y) not in done:
+                        n.append((cost+1, (x,y)))
+        check(x-1, y)
+        check(x+1,y)
+        check(x, y-1)
+        check(x, y+1)
+
+        for p in n:
+            heapq.heappush(q, p)
+
+solve(1)
+solve(2)