advent_of_code_2021

day11_part1.py (3830B)

---

 #!/usr/bin/env python
 """Advent of Code 2021, day 11 (part 1): Flashing octopuses
 For part 1, we're simulating bioluminescent octopuses that 'flash' when they
 reach a threshold 'energy level'."""
 
 # Algorithm for simulating one 'step' of the octopuses (this is basically
 # breadth-first search):
 # initialize "increments" to 1
 # initialize "already flashed" to false
 # while any increments are 1
 #   add increments to energy
 #   set increments to 0
 #   set flashing now to energy > 9 and not already flashed
 #   set already flashed to already flashed or flashing now
 #   for flashing octopus in flashing now
 #     increase neighbor's increments by 1
 # Note: these functions kinda ignore Python's pass-by-reference behavior,
 # they're updating numpy arrays in place and also returning them, but that's
 # not causing any problems here.
 
 from typing import Tuple
 import itertools
 import functools
 import numpy as np
 from utils import get_puzzle_input, convert_input_to_array
 
 EXAMPLE_INPUT = \
 """5483143223
 2745854711
 5264556173
 6141336146
 6357385478
 4167524645
 2176841721
 6882881134
 4846848554
 5283751526
 """
 
 def increment_neighbors(increments: np.ndarray,
                         row: int, col: int) -> np.ndarray:
     """Increase the value of the given `increments` 2d array by 1 in the
     8-connected neighborhood of the element at the given row and column"""
     # Note that this doesn't check if `row` and `col` are valid, but it
     # will fail silently (possibly updating elements at the edges
     for row_offset, col_offset in itertools.product(range(-1, 2), range(-1, 2)):
         if row_offset != 0 or col_offset != 0:
             row_idx = row + row_offset
             col_idx = col + col_offset
             if 0 <= row_idx < increments.shape[0] and \
                     0 <= col_idx < increments.shape[1]:
                 increments[row_idx, col_idx] += 1
     return increments
 
 def simulate_step(energy_levels: np.ndarray,
                   flashes: int = 0) -> Tuple[np.ndarray, int]:
     """For the given 2d array of energy levels and the number of flashes
     detected so far, simulate one step of the flashing octopuses and return a
     tuple containing the new energy levels and an updated number of flashes"""
     increments = np.ones_like(energy_levels)
     already_flashed = np.zeros_like(energy_levels, dtype=bool)
 
     while increments.any():
         # Apply the current set of increments
         energy_levels += increments
         increments[:,:] = 0
         # Figure out who's flashing now and update who's flashed so far
         flashing_now = (energy_levels > 9) & (~already_flashed)
         already_flashed |= flashing_now
         # Update neighbors' energy levels (i.e., determine who will be
         # incremented on the next step through the loop)
         for row, col in zip(*np.where(flashing_now)):
             increments = increment_neighbors(increments, row, col)
 
     # Reset the flashing octopuses' energy levels to 0
     energy_levels[already_flashed] = 0
     return (energy_levels,
             flashes + already_flashed.sum())
 
 def run_simulation(energy_levels: np.ndarray, steps: int) -> int:
     """Run the simulation from the initial 2d array of energy levels for the
     given number of steps, and return the number of flashes detected"""
     return functools.reduce(
         lambda x, y: simulate_step(*x),
         range(steps),
         (energy_levels, 0)
     )[1]
 
 def solve_puzzle(input_string: str) -> int:
     """Return the numeric solution to the puzzle"""
     return run_simulation(convert_input_to_array(input_string), 100)
 
 def main() -> None:
     """Run when the file is called as a script"""
     assert solve_puzzle(EXAMPLE_INPUT) == 1656
     print("Number of flashes after 100 steps:",
           solve_puzzle(get_puzzle_input(11)))
 
 if __name__ == "__main__":
     main()