diff --git a/2021/d11/ex2/ex2.py b/2021/d11/ex2/ex2.py new file mode 100755 index 0000000..0dee9eb --- /dev/null +++ b/2021/d11/ex2/ex2.py @@ -0,0 +1,67 @@ +#!/usr/bin/env python + +import itertools +import sys +from copy import deepcopy +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()