Takodachi! Advent Of Code 2021, Day 11

#!/usr/bin/env python3

import numpy as np
from functools import reduce

input = """1224346384
5621128587
6388426546
1556247756
1451811573
1832388122
2748545647
2582877432
3185643871
2224876627"""


def flash(iter, takodachi, count, all_flash, flash_threshold=10):
    takodachi += 1
    width, height = takodachi.shape
    old = set()
    while True:
        flash_mask = np.where(np.greater_equal(takodachi, flash_threshold))
        d = list(set([(flash_mask[0][i], flash_mask[1][i]) for i in range(len(flash_mask[0]))]).difference(old))
        if len(d) > 0:
            mask = (np.array([d[i][0] for i in range(len(d))]), np.array([d[i][1] for i in range(len(d))]))
            old = old.union(set([(mask[0][i], mask[1][i]) for i in range(len(mask[0]))]))
            adjacents = []
            for x in [-1, 0, 1]:
                for y in [-1, 0, 1]:
                    if x != 0 or y != 0:
                        mx = mask[0] + x
                        my = mask[1] + y
                        keep = np.where((mx >= 0) & (mx < width) & (my >= 0) & (my < width))[0]
                        adjacents.append((mx[keep], my[keep]))
            for adjacents_mask in adjacents:
                takodachi[adjacents_mask] += 1
        else:
            break

    d = list(old)
    flashes = len(d)
    mask = (np.array([d[i][0] for i in range(flashes)]), np.array([d[i][1] for i in range(flashes)]))
    if flashes > 0:
        takodachi[mask] = 0
        if flashes == width * height:
            all_flash.append(iter + 1)
    return takodachi, flashes + count, all_flash


if __name__ == '__main__':
    width = len(input.split("\n", 1)[0])
    n = np.array([c for c in input.replace('\n', '')], dtype=int)
    tako = n.reshape(width, -1)
    max_step = 500
    tako, flashes, all_flashes = reduce(lambda acc, iter: flash(iter, acc[0], acc[1], acc[2]), np.arange(max_step), (tako, 0, []))
    print(tako, flashes, all_flashes[0])