adventofcode2022/day12/sol.py

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]

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')
            start = (x,y)
        elif get(x,y) == 0: # 'E'
            data[y][x] = chars.index('z')
            goal = (x,y)
        elif get(x,y) == 2: # 'a'
            # for part 2
            extra.append((x,y))

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

        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)