day 12 python simplify

day12/sol.py

@@ -1,28 +1,27 @@
 import heapq
 import string
 
-chars = 'ES'+string.ascii_lowercase
+elevation = {c:i for i, c in enumerate(string.ascii_lowercase)}
+elevation['E'] = elevation['z']
+elevation['S'] = elevation['a']
 
 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]
+#print(*data, sep="\n")
 
 extra = []
-for x in range(len(data[0])):
-    for y in range(len(data)):
-        if get(x,y) == 1: # 'S'
-            data[y][x] = chars.index('a')
+for y, row in enumerate(data):
+    for x, c in enumerate(row):
+        if c == 'S':
             start = (x,y)
-        elif get(x,y) == 0: # 'E'
-            data[y][x] = chars.index('z')
+        elif c == 'E':
             goal = (x,y)
-        elif get(x,y) == 2: # 'a'
+        elif c == 'a':
             # for part 2
             extra.append((x,y))
 
+def get(x, y):
+    return elevation[data[y][x]]
+
 def solve(part):
     q = [(0, start)]
 
@@ -41,19 +40,15 @@ def solve(part):
         here = get(x,y)
         done[(x,y)] = cost
 
-        n = []
-        def check(x,y):
+        def push(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)
+                        heapq.heappush(q, (cost+1, (x,y)))
+        push(x-1, y)
+        push(x+1,y)
+        push(x, y-1)
+        push(x, y+1)
 
 solve(1)
 solve(2)