adventofcode2022/day12/sol.py

55 lines
1.2 KiB
Python

import heapq
import string
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)
#print(*data, sep="\n")
extra = []
for y, row in enumerate(data):
for x, c in enumerate(row):
if c == 'S':
start = (x,y)
elif c == 'E':
goal = (x,y)
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)]
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
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:
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)