advent/y2022/p12.py

def parse_grid(text):
    grid = {}
    for y, row in enumerate(text):
        for x, v in enumerate(row):
            grid[complex(x, y)] = v if v in 'SE' else ord(v)
    return grid


def print_grid(mapping):
    xmin, *_, xmax = {int(p.real) for p in mapping}
    ymin, *_, ymax = {int(p.imag) for p in mapping}
    for y in range(ymin, ymax + 1):
        for x in range(xmin, xmax + 1):
            print(mapping.get(complex(x, y), ' '), end=' ')
        print()


def shortest_path(starts, ends, evolve):
    starts = {k: [k] for k in starts}
    known = starts.copy()
    while starts:
        old = list(starts)[0]
        path = starts.pop(old)
        starts |= {new: path + [new] for new in evolve(old) if new not in known}
        known |= starts
        for final in set(starts) & set(ends):
            # print_grid({k: '.' for k in known[final]})
            return known[final]


def main():
    text = open(0).read().splitlines()
    grid = parse_grid(text)

    rev = {v: k for k, v in grid.items()}
    start, end = rev['S'], rev['E']
    grid[start], grid[end] = ord('a'), ord('z')

    dof = [1, -1, 1j, -1j]
    evolve = lambda old: {old + step for step in dof if is_valid(old, step)}
    is_valid = lambda old, step: old + step in grid and grid[old + step] - grid[old] <= 1

    print(len(shortest_path({start}, {end}, evolve)) - 1)
    starts = {k for k, v in grid.items() if v == ord('a')}
    print(len(shortest_path(starts, {end}, evolve)) - 1)


main()