day 12 python tweaks

day12/sol.py

@@ -1,46 +1,51 @@
-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:
+        elif get(x,y) == 2: # 'a'
-            # part 2
+            # for part 2
-            q.append((0,(x,y)))
+            extra.append((x,y))
-            pass
 
-q.append((0, start))
+def solve(part):
-path = {}
+    q = [(0, start)]
-while q:
+
+    if part == 2:
+        q.extend((0,p) for p in extra)
+
+    done = {}
+    while q:
         cost, (x,y) = heapq.heappop(q)
-    if (x,y) in path:
+        if (x,y) in done:
             continue
         #print(x,y,cost)
         if (x,y) == goal:
             print('found', cost)
             break
         here = get(x,y)
-    path[(x,y)] = cost
+        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 path:
+                    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)
@@ -50,3 +55,5 @@ while q:
         for p in n:
             heapq.heappush(q, p)
 
+solve(1)
+solve(2)