width = 10 height = 10 grid = {} with open("input.txt") as f: lines = f.read().split('\n') for i in range(height): for j in range(width): grid[(i,j)] = int(lines[i][j]) def steps(grid, steps): """ Step through grid multiple times and count the flashes. """ flashes = 0 for i in range(steps): flashes += step(grid) return flashes def step(grid): """ Step through grid once and count the flashes. """ increment(grid) flashes = flash(grid) return flashes def increment(grid): """ Increment the entire grid. """ for i in range(height): for j in range(width): grid[(i,j)] += 1 def flash(grid): """ Calculate which octopi should flash. Also count how many flashes. """ flashes = 0 flashed = set() cascade = [] # find which octopi should flash and add them to queue for i in range(height): for j in range(width): point = (i,j) if grid[point] > 9 and point not in flashed: flashed.add(point) cascade += neighbors(point) grid[point] = 0 # reset flashes += 1 # process the queue of flashing octopi while cascade: point = cascade.pop() # if it hasn't already flashed, increment its energy if point not in flashed: grid[point] = grid[point] + 1 # then figure out if it should flash if grid[point] > 9 and point not in flashed: flashed.add(point) cascade += neighbors(point) grid[point] = 0 # reset flashes += 1 return flashes def neighbors(point): """ Get a list of neighboring points, diagonals included. """ i,j = point n = [ (i-1,j-1), (i-1,j), (i-1,j+1), (i,j-1), #(i,j) is excluded (i,j+1), (i+1,j-1), (i+1,j), (i+1,j+1) ] return [ (i,j) for (i,j) in n if (-1 < i < height) and (-1 < j < width) ] def display(grid): """ Pretty print the grid. """ for i in range(height): row = [] for j in range(width): row.append(grid[(i,j)]) print(row) grid1 = grid.copy() part1 = steps(grid1,100) #display(grid1) print(part1) grid2 = grid.copy() from collections import Counter i = 0 while True: step(grid2) i += 1 c = Counter(grid2.values()) if c[0] == width * height: # all are flashing at same time break print(i)