import argparse import functools from typing import NamedTuple, List, Dict, Set, Tuple 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 return octopi_list octs = octopi for x in range(0, 100): octs = incrementOctopi(octs) octs = flashOctopi(octs) print(TOTAL_FLASHES)