advent_of_code_2021/day11/part2.py

import argparse
import functools
from typing import NamedTuple, List, Dict, Set, Tuple
import sys

parser = argparse.ArgumentParser()
parser.add_argument("ifile", type=argparse.FileType('r'))
args = parser.parse_args()

octopi = [[int(x) for x in line.strip()] for line in args.ifile.readlines()]
max_dim = len(octopi) # thing is square, so this works for both
TOTAL_FLASHES = 0


def getNeighbors(x: int, y: int):
    positions = [(x - 1, y - 1), (x, y - 1), (x + 1, y - 1),
                 (x - 1, y), (x, y),  (x + 1, y),
                 (x - 1, y + 1), (x, y + 1), (x + 1, y + 1)]
    return list(filter(lambda pair: 0 <= pair[0] < max_dim and 0 <= pair[1] < max_dim, positions))


def printOctopi(octopi_list: List[List[int]]):
    for line in octopi_list:
        print("".join([str(x) for x in line]))


def incrementOctopi(octopi_list: List[List[int]]) -> List[List[int]]:
    return [list(map(lambda x: x + 1, line)) for line in octopi_list]


def flashOctopi(octopi_list: List[List[int]]) -> List[List[int]]:
    global TOTAL_FLASHES
    to_flash: Set[Tuple[int, int]] = set()
    has_flashed: Set[Tuple[int, int]] = set()

    for y in range(0, max_dim):
        for x in range(0, max_dim):
            if octopi_list[y][x] > 9:
                to_flash.add((x, y))

    while(len(to_flash) > 0):
        octopode = to_flash.pop()
        TOTAL_FLASHES += 1
        has_flashed.add(octopode)

        for oct_pos in getNeighbors(octopode[0], octopode[1]):
            octopi_list[oct_pos[1]][oct_pos[0]] += 1
            if octopi_list[oct_pos[1]][oct_pos[0]] > 9:
                if (oct_pos[0], oct_pos[1]) not in has_flashed:
                    to_flash.add((oct_pos[0], oct_pos[1]))

    for pos in has_flashed:
        octopi_list[pos[1]][pos[0]] = 0

    if len(has_flashed) == 100:
        raise Exception

    return octopi_list

octs = octopi
for x in range(0, 1000):
    octs = incrementOctopi(octs)
    try:
        octs = flashOctopi(octs)
    except:
        print(x + 1)
        break