advent/y2022/p24.py

def loop(p, drx):
    p += drx
    if p not in grid:
        p -= drx * (H if drx.imag else W)
    return p


def blizzards(n, cache={}, valid={'>': 1, '<': -1, '^': -1j, 'v': 1j}):
    if n not in cache:
        if n == 0:
            cache[n] = [(p, valid[c]) for p, c in grid.items() if c != '.']
        else:
            cache[n] = [(loop(p, drx), drx) for p, drx in blizzards(n - 1)]
    return cache[n]


def bfs(starts, ends, tt=0, dir5=[0, 1, -1j, -1, 1j]):
    edge = starts
    while not ends & edge:
        tt += 1
        blizz = {p for p, _ in blizzards(tt)}
        edge = {p + dp for p in edge for dp in dir5} & set(grid) - blizz
    return tt


def main():
    global grid, W, H

    rows = open(0).read().splitlines()
    grid = {complex(x, y): c for y, r in enumerate(rows)
                             for x, c in enumerate(r) if c != '#'}

    H, W = len(rows) - 2, len(rows[1]) - 2
    start, *_, end = sorted(grid, key=lambda p: (p.imag, p.real))

    t1 = bfs({start}, {end})
    print(t1)

    t2 = bfs({end}, {start}, t1)
    t3 = bfs({start}, {end}, t2)
    print(t3)


main()