AoC 2022 Day 24 - Python

def parse(data: str):
    rows = data.split("\n")
    bottom_wall = len(rows) - 2
    right_wall = len(rows[0]) - 2
    up, down, left, right = [], [], [], []
    for i, row in enumerate(rows):
        for j, spot in enumerate(row):
            if spot == '^':
                up.append(complex(j-1, i-1))
            elif spot == 'v':
                down.append(complex(j-1, i-1))
            elif spot == '<':
                left.append(complex(j-1, i-1))
            elif spot == '>':
                right.append(complex(j-1, i-1))
    return up, down, left, right, bottom_wall, right_wall

def move_blizzard(up: list[complex], down: list[complex], left: list[complex], right: list[complex], right_wall: int, bottom_wall: int):
    for i, b in enumerate(up):
        b -= 1j
        if b.imag < 0:
            b += complex(0, bottom_wall)
        up[i] = b
    for i, b in enumerate(down):
        b += 1j
        if b.imag >= bottom_wall:
            b = b.real
        down[i] = b
    for i, b in enumerate(left):
        b -= 1
        if b.real < 0:
            b += right_wall
        left[i] = b
    for i, b in enumerate(right):
        b += 1
        if b.real >= right_wall:
            b -= right_wall
        right[i] = b

def print_blizzard(up: list[complex], down: list[complex], left: list[complex], right: list[complex], right_wall: int, bottom_wall: int, queue: list[complex]):
    # print_blizzard(up, down, left, right, right_wall, bottom_wall)
    out = [['.']*right_wall for _ in range(bottom_wall)]
    for b in up:
        out[int(b.imag)][int(b.real)] = '^'
    for b in down:
        out[int(b.imag)][int(b.real)] = 'v'
    for b in left:
        out[int(b.imag)][int(b.real)] = '<'
    for b in right:
        out[int(b.imag)][int(b.real)] = '>'
    out = [['#'] + row + ['#'] for row in out]
    out.insert(0, ['#', '.'] + ['#']*right_wall)
    out.append(['#']*right_wall + ['.', '#'])
    for you in queue:
        out[int(you.imag)+1][int(you.real)+1] = 'E'
    for row in out:
        print(''.join(row))

def blizzard(data: str):
    up, down, left, right, bottom_wall, right_wall = parse(data)
    queue = [-1j]
    moves = 0
    destinations = [-1j, complex(right_wall-1, bottom_wall)]
    destination = complex(right_wall-1, bottom_wall)
    wall = set()
    for i in range(-1, right_wall+1):
        wall.add(complex(i, -1))
        wall.add(complex(i, bottom_wall))
    for i in range(-1, bottom_wall+1):
        wall.add(complex(-1, i))
        wall.add(complex(right_wall, i))
    wall.add(-2j)
    wall.add(complex(right_wall-1, bottom_wall+1))
    wall -= set([-1j, complex(right_wall-1, bottom_wall)])
    while True:
        move_blizzard(up, down, left, right, right_wall, bottom_wall)
        next_queue = set()
        moves += 1
        blizzards = set(up+down+left+right).union(wall)
        while len(queue):
            you = queue.pop(0)
            if you == destination:
                print(moves - 1)
                if len(destinations):
                    destination = destinations.pop(0)
                    queue = []
                    next_queue = set()
                else:
                    return
            for i in [0, 1, -1, 1j, -1j]:
                move = you + i
                if move not in blizzards:
                    next_queue.add(move)
        queue = list(next_queue)

import time
st = time.time()

# blizzard(example)
blizzard(input)
print(time.time() - st)