Towards world domINAtion!
#!/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])
