AoC 2021 Day 11

import numpy as np
import os


ALL_FLASHING = np.zeros((10, 10), dtype=np.int32)


def read_and_init(fn):
    with open(fn, 'r') as f:
        lines = f.read().splitlines()
    grid = [[int(x) for x in line] for line in lines]
    oktapodi = np.array(grid, dtype=np.int32)
    return oktapodi


def run(octopi: np.ndarray, part: int):
    flashes = 0
    i = 0
    while True:
        octopi += 1
        flashed = set()
        f = np.where(octopi == 10)
        flashing = [(f[0][i], f[1][i]) for i in range(len(f[0]))]
        while flashing:
            r, c = flashing.pop(0)
            flashed.add((r, c))
            r0, r1 = max(r-1, 0), max(r+2, 0)
            c0, c1 = max(c-1, 0), max(c+2, 0)
            neighbors = octopi[r0:r1, c0:c1]
            neighbors += 1
            f = np.where(neighbors == 10)
            flashing_neighbors = [(f[0][i] + r0, f[1][i] + c0) for i in range(len(f[0]))]
            flashing += flashing_neighbors
        for r, c in flashed:
            octopi[r, c] = 0
        flashes += octopi.size - np.count_nonzero(octopi)
        i += 1
        if part == 1 and i == 100:
            print('Part 1:', flashes)
            break
        if part == 2:
            export_image(octopi, i)
            if np.all(octopi == ALL_FLASHING):
                print('Part 2: ', i)
        if i == 500:
            break


oktos = read_and_init(os.path.dirname(__file__) + '\\in.txt')
run(oktos, 1)
oktos = read_and_init(os.path.dirname(__file__) + '\\in.txt')
run(oktos, 2)