wine

y2021/p11.py

+import itertools
+import toolkit
+
+
+state = {k: int(v) for k, v in toolkit.read_image(text)[0].items()}
+D8 = [dx + dy for dx in [1, 0, -1] for dy in [1j, 0, -1j] if dx + dy]
+
+
+def flash(k, flashed_this_round):
+    state[k] = 0
+    for p in (k + step for step in D8):
+        if p in state and p not in flashed_this_round:
+            state[p] += 1
+
+
+ans1 = 0
+for i in itertools.count(1):
+    for k in state:
+        state[k] += 1
+
+    flashed_this_round = set()
+    while any(v > 9 for v in state.values()):
+        for k, v in state.items():
+            if v > 9:
+                flash(k, flashed_this_round)
+                flashed_this_round.add(k)
+
+    if i <= 100:
+        ans1 += len(flashed_this_round)
+
+    if len({v for v in state.values()}) == 1:
+        ans2 = i
+        break