Advent of Code 2020 Day 17

AoC2020_day17.py

def build(initial, dim_sz=3):
    cubes = {}
    for y, line in enumerate(initial):
        for x, state in enumerate(line):
            cubes[(x, y, 0, 0)] = state
    fill(cubes, dim_sz=dim_sz)
    return cubes

def should_flip(cubes, pos, dim_sz=3):
    DW_MOV = (0,) if dim_sz == 3 else (-1, 0, 1)

    state = cubes[pos]
    px, py, pz, pw = pos

    n = 0
    for dx in (-1, 0, 1):
        for dy in (-1, 0, 1):
            for dz in (-1, 0, 1):
                for dw in DW_MOV:
                    if dx == 0 and dy == 0 and dz == 0 and dw == 0:
                        continue
                    if cubes.get((px+dx, py+dy, pz+dz, pw+dw), ".") == "#":
                        n += 1
                    if n > 3:
                        break

    if (state == "#" and n not in (2, 3)) or (state == "." and n == 3):
        return True

    return False


def fill(cubes, subset=None, dim_sz=3):
    DW_MOV = (0,) if dim_sz == 3 else (-1, 0, 1)
    if not subset:
        subset = list(cubes.keys())
    for px, py, pz, pw in subset:
        for dx in (-1, 0, 1):
            for dy in (-1, 0, 1):
                for dz in (-1, 0, 1):
                    for dw in DW_MOV:
                        if (px+dx, py+dy, pz+dz, pw+dw) not in cubes:
                            cubes[(px+dx, py+dy, pz+dz, pw+dw)] = "."

N_CYCLE = 6
DIM_SZ = 4  # Change to 4 for Part 2

cubes = build([line.strip() for line in
                open("input_17").readlines()], DIM_SZ)
for i in range(N_CYCLE):
    to_flip = []
    for pos in cubes.keys():
        if should_flip(cubes, pos, DIM_SZ):
            to_flip.append(pos)
    for pos in to_flip:
        if cubes[pos] == "#":
            cubes[pos] = "."
        else:
            cubes[pos] = "#"
    fill(cubes, to_flip, DIM_SZ)
    print(i+1, list(cubes.values()).count("#"))