aoc-2022/24.py

from queue import PriorityQueue

lines = [line.strip() for line in open("24.input")]

blizzards = set()

size = (len(lines[0]), len(lines))

DIRS = {
        "<": (-1, 0),
        ">": (1, 0),
        "v": (0, 1),
        "^": (0, -1),
        }

for y, line in enumerate(lines):
    for x, c in enumerate(line):
        if c in ("<", ">", "v", "^"):
            blizzards.add((x, y, c))

def show(grid):
    for y in range(size[1]):
        for x in range(size[0]):
            if x == 0 or y == 0 or x == size[0] - 1 or y == size[1] - 1:
                print("#", end="")
                continue
            blizz = [c for (x_, y_, c) in grid if x == x_ and y == y_]
            if len(blizz) > 1:
                print(len(blizz), end="")
            elif len(blizz) == 1:
                print(blizz[0], end="")
            else:
                print(".", end="")
        print("")
    print("-" * 20)

def in_bound(pos):

    start = (1, 0)
    end = (size[0] - 2, size[1] - 1)
    if pos in [start, end]:
        return True

    nx, ny = pos
    if nx <= 0:
        return False
    if nx >= size[0] - 1:
        return False
    if ny <= 0:
        return False
    if ny >= size[1] - 1:
        return False
    return True

def next_map(current):
    new = set()
    for (x, y, d) in current:
        dx, dy = DIRS[d]
        nx, ny = x + dx, y + dy
        if nx == 0:
            nx = size[0] - 2
        if nx == size[0] - 1:
            nx = 1
        if ny == 0:
            ny = size[1] - 2
        if ny == size[1] - 1:
            ny = 1
        new.add((nx, ny, d))
    return new

def search(grid_cache, runs):

    start = (1, 0)
    end = (size[0] - 2, size[1] - 1)

    q = PriorityQueue()
    q.put((0, start, 0))

    dist_cache = {(0, start, 0): 0}

    while not q.empty():
        tick, pos, run = q.get()
        #print(tick, pos, run)

        nxt_map = grid_cache[(tick + 1) % len(grid_cache)]

        if pos == start and run == 1:
            q.put((tick, pos, 2))

        if pos == end:
            if run == runs:
                return tick
            elif run == 0:
                q.put((tick, pos, run + 1))

        for dx, dy in [(1, 0), (0, 1), (-1, 0), (0, -1), (0, 0)]:
            dest = (pos[0] + dx, pos[1] + dy)
            if not in_bound(dest):
                continue

            if dest not in nxt_map:
                new_state = ((tick + 1) % len(grid_cache), dest, run)
                if dist_cache.get(new_state, 100000000) > tick + 1:
                    dist_cache[new_state] = tick + 1
                    q.put((tick + 1, dest, run))


grid_cache = []

current = blizzards

while True:
    grid_cache.append(current)
    current = next_map(current)
    if current in grid_cache:
        break

grid_cache = [set((x, y) for (x, y, c) in current) for current in grid_cache]

print(search(grid_cache, 0))
print(search(grid_cache, 2))