aoc-2020/17.py

grid_3d = {}
grid_4d = {}

with open("17.input") as f:
    for (y, line) in enumerate(f.readlines()):
        for (x, c) in enumerate(list(line.strip())):
            grid_3d[(x, y, 0)] = c
            grid_4d[(x, y, 0, 0)] = c

def get_surrounding(grid, locked_coords, free_coords):
    if len(free_coords) == 0:
        if grid.get(locked_coords, ".") == "#":
            return 1
        else:
            return 0

    count = 0
    for x in range(free_coords[0] - 1, free_coords[0] + 2):
        count += get_surrounding(grid, tuple(list(locked_coords) + [x]), free_coords[1:])

    # Remove self
    if len(locked_coords) == 0 and grid.get(free_coords, '.') == '#':
        count -= 1

    return count


def calc_new(grid, coords):
    count = get_surrounding(grid, (), coords)
    if grid.get(coords, '.') == '#':
        if count in [2, 3]:
            return '#'
        else:
            return '.'
    else:
        if count == 3:
            return '#'
        else:
            return '.'


def cycle(grid, fixed_coords, sizes):
    new_grid = {}
    if sizes == []:
        new_grid[tuple(fixed_coords)] = calc_new(grid, tuple(fixed_coords))
    else:
        for x in sizes[0]:
            new_grid.update(cycle(grid, fixed_coords + [x], sizes[1:]))

    return new_grid


def cycle_multiple(grid, start_sizes):
    my_grid = dict(grid)
    for i in range(1, 7):
        sizes = [range(-i, s + i) for s in start_sizes]
        my_grid = cycle(my_grid, [], sizes)
    return my_grid


print("Answer 1:", len([1 for c in cycle_multiple(grid_3d, (8, 8, 1)).values() if c == '#']))
print("Answer 2:", len([1 for c in cycle_multiple(grid_4d, (8, 8, 1, 1)).values() if c == '#']))