
# Python imports
from pathlib import Path
from typing import Callable, Dict, Generator, Iterable, List, Tuple

import numpy as np

# Paths to data
testpath = Path("day11_test.txt")
flashtestpath = Path("day11_flashtest.txt")
datapath = Path("day11_data.txt")


def load_input(fpath: Path) -> np.array:
    """Return octopus states as an array of ints
    
    :param fpath:  Path to data file
    """
    return np.genfromtxt(fpath, dtype=int, delimiter=1)


def update_neighbours(arr: np.array, row: int, col: int) -> np.array:
    """Returns array with all neighbours at (row, col) incremented
    
    :param arr:  array of octopus states
    :param row:  row of flashing octopus
    :param col:  col of flashing octopus
    """
    padded = np.pad(arr, 1)
    row, col = row+1, col+1
    padded[row-1:row+2, col-1:col+2] += 1
    return padded[1:-1, 1:-1]

def get_octopus_flashcount(arr: np.array, steps: int) -> Tuple[np.array, int]:
    """Returns octopus energy array and count of octopus flashes
    
    :param arr:  octopus energy state array
    :param steps:  number of steps to run
    """
    flashcount = 0  # tally of octopuses that flashed
    
    # Iterate over the required number of steps
    for step in range(steps):
        arr += 1  # increment elements in array
        flashmask = arr > 9  # mask of octopuses that flash
        flashed = flashmask  # mask of octupuses that flashed in this step
        
        # It is possible that the first lot of flashing octopuses can
        # cause other octopuses to flash, so we need to iterate until
        # the cascade of flashing stops
        while flashmask.sum():
            flashcount += flashmask.sum()  # count octopuses currently flashing
            rows, cols = np.where(flashmask)  # get locations of flashing octopuses
            for row, col in list(zip(rows, cols)):  # update status array
                arr = update_neighbours(arr, row, col)
            arr[flashed] = 0  # reset octopuses that originally flashed
            flashmask = arr > 9  # mask of octopuses that now flash
            flashed += flashmask  # update mask of octopuses that flashed in this step

    return arr, flashcount


data = load_input(flashtestpath)
get_octopus_flashcount(data, 2)

(array([[4, 5, 6, 5, 4],
        [5, 1, 1, 1, 5],
        [6, 1, 1, 1, 6],
        [5, 1, 1, 1, 5],
        [4, 5, 6, 5, 4]]),
 9)


data = load_input(testpath)
get_octopus_flashcount(data, 100)

(array([[0, 3, 9, 7, 6, 6, 6, 8, 6, 6],
        [0, 7, 4, 9, 7, 6, 6, 9, 1, 8],
        [0, 0, 5, 3, 9, 7, 6, 9, 3, 3],
        [0, 0, 0, 4, 2, 9, 7, 8, 2, 2],
        [0, 0, 0, 4, 2, 2, 9, 8, 9, 2],
        [0, 0, 5, 3, 2, 2, 2, 8, 7, 7],
        [0, 5, 3, 2, 2, 2, 2, 9, 6, 6],
        [9, 3, 2, 2, 2, 2, 8, 9, 6, 6],
        [7, 9, 2, 2, 2, 8, 6, 8, 6, 6],
        [6, 7, 8, 9, 9, 9, 8, 7, 6, 6]]),
 1656)


data = load_input(datapath)
get_octopus_flashcount(data, 100)

(array([[5, 5, 6, 8, 4, 5, 5, 9, 7, 6],
        [5, 6, 8, 6, 4, 0, 0, 0, 9, 7],
        [6, 8, 5, 8, 0, 0, 0, 0, 0, 9],
        [8, 7, 2, 4, 0, 0, 0, 0, 0, 0],
        [1, 4, 2, 3, 5, 0, 0, 0, 0, 0],
        [1, 5, 2, 2, 3, 4, 5, 9, 8, 0],
        [2, 5, 2, 2, 2, 2, 2, 7, 5, 5],
        [0, 6, 3, 2, 2, 2, 2, 7, 4, 4],
        [0, 0, 5, 3, 2, 2, 2, 6, 4, 4],
        [0, 0, 0, 3, 2, 2, 7, 5, 4, 4]]),
 1743)


def get_octopus_sync(arr: np.array) -> Tuple[np.array, int]:
    """Returns octopus energy array and steps to convergence
    
    :param arr:  octopus energy state array
    :param steps:  number of steps to run
    """    
    step = 0  # Count of steps taken
    
    # Iterate until all values in the status array are zero
    while arr.sum():
        arr += 1  # increment elements in array
        flashmask = arr > 9  # mask of octopuses that flash
        flashed = flashmask  # mask of octupuses that flashed in this step
        
        # It is possible that the first lot of flashing octopuses can
        # cause other octopuses to flash, so we need to iterate until
        # the cascade of flashing stops
        while flashmask.sum():
            rows, cols = np.where(flashmask)  # get locations of flashing octopuses
            for row, col in list(zip(rows, cols)):  # update status array
                arr = update_neighbours(arr, row, col)
            arr[flashed] = 0  # reset octopuses that originally flashed
            flashmask = arr > 9  # mask of octopuses that now flash
            flashed += flashmask  # update mask of octopuses that flashed in this step
        step += 1

    return arr, step


data = load_input(testpath)
get_octopus_sync(data)

(array([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]),
 195)


data = load_input(datapath)
get_octopus_sync(data)

(array([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]),
 364)

Day 11: Dumbo Octopus¶

Puzzle 2:¶