day11_2.py

day11_2.py

def within_bounds(_y, _x, matrix):
    return 0 <= _y < len(matrix) and 0 <= _x < len(matrix[0])


def adjacent_cells(y, x, matrix):
    for _y in range(y - 1, y + 2):
        for _x in range(x - 1, x + 2):
            if within_bounds(_y, _x, matrix):
                yield _y, _x


def cells(matrix):
    for y in range(len(matrix)):
        for x in range(len(matrix[0])):
            yield y, x


def increase_energy_level(y, x, matrix, path_matrix):
    flash_count = 0
    if path_matrix[y][x] != 1:
        matrix[y][x] += 1
        if matrix[y][x] == 10:
            path_matrix[y][x] = 1
            matrix[y][x] = 0
            flash_count += 1 + sum(increase_energy_level(_y, _x, matrix, path_matrix)
                                   for _y, _x in adjacent_cells(y, x, matrix))
    return flash_count


if __name__ == "__main__":
    with open("input3.txt") as file:
        lines = [line.strip() for line in file.readlines()]

    matrix = []
    path_matrix = []
    for line in lines:
        matrix.append(list(map(int, [c for c in line])))
        path_matrix.append([0] * len(line))

    step_counter = 1
    while True:
        flash_count = sum(increase_energy_level(y, x, matrix, path_matrix) for y, x in cells(matrix))
        if flash_count == len(matrix) * len(matrix[0]):
            break
        for y, x in cells(matrix):
            path_matrix[y][x] = 0
        step_counter += 1

    print(step_counter)
    assert step_counter == 515