From 00e984420cbd46c464d5fac7b7a7f87a15afe429 Mon Sep 17 00:00:00 2001
From: Bruno BELANYI <bruno@belanyi.fr>
Date: Sat, 11 Dec 2021 14:12:19 +0100
Subject: [PATCH] 2021: d11: ex2: add solution

---
 2021/d11/ex2/ex2.py | 66 +++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 66 insertions(+)
 create mode 100755 2021/d11/ex2/ex2.py

diff --git a/2021/d11/ex2/ex2.py b/2021/d11/ex2/ex2.py
new file mode 100755
index 0000000..39e5092
--- /dev/null
+++ b/2021/d11/ex2/ex2.py
@@ -0,0 +1,66 @@
+#!/usr/bin/env python
+
+import itertools
+import sys
+from typing import Iterator, List, Set, Tuple
+
+Grid = List[List[int]]
+Point = Tuple[int, int]
+
+
+def solve(input: List[str]) -> int:
+    levels = [[int(c) for c in line] for line in input]
+
+    def step(levels: Grid) -> Tuple[Grid, int]:
+        # First step, increase levels
+        levels = [[l + 1 for l in line] for line in levels]
+
+        def excited(levels: Grid) -> Set[Point]:
+            return set(
+                (i, j)
+                for i in range(len(levels))
+                for j in range(len(levels[i]))
+                if levels[i][j] > 9
+            )
+
+        def neighbours_of(point: Point) -> Iterator[Point]:
+            for dx, dy in itertools.product((-1, 0, 1), repeat=2):
+                if dx == 0 and dy == 0:
+                    continue
+                x = point[0] + dx
+                y = point[1] + dy
+                if x < 0 or x >= len(levels):
+                    continue
+                if y < 0 or y >= len(levels[x]):
+                    continue
+                yield x, y
+
+        # Second step, do flashes
+        has_flashed: Set[Point] = set()
+        while len(flashes := (excited(levels) - has_flashed)) > 0:
+            for (i, j) in flashes:
+                has_flashed.add((i, j))
+                for x, y in neighbours_of((i, j)):
+                    levels[x][y] += 1
+
+        # Finally, bring back energy levels to 0
+        for i, j in has_flashed:
+            levels[i][j] = 0
+
+        return levels, len(has_flashed)
+
+    for i in itertools.count(1):
+        levels, flashes = step(levels)
+        if flashes == len(list(itertools.chain.from_iterable(levels))):
+            return i
+
+    assert False  # Sanity check
+
+
+def main() -> None:
+    input = [line.strip() for line in sys.stdin.readlines()]
+    print(solve(input))
+
+
+if __name__ == "__main__":
+    main()