conway cubes

approach

active cubes are a set of coordinate tuples. the 2d input is padded with zeros up to the requested number of dimensions, and every function after that works on tuples of any length: neighbors takes the product of (0, -1, 1) over each axis and skips the first combination, which is all zeros.

each step counts, for every cell next to an active cube, how many active neighbours it has. an active cube stays active with 2 or 3, and any cell with exactly 3 becomes active. part 1 runs 6 steps in 3d and part 2 the same 6 steps in 4d; the only difference is the dimension passed to parse.

the trick

only store active cells and count neighbours outwards from them, and the same code works in any number of dimensions.

def step(grid):
    counts, ngrid = defaultdict(int), set()
    for p in grid:
        for np in neighbors(p):
            counts[np] += 1
    for p, v in counts.items():
        if p in grid and v in [2, 3]:
            ngrid.add(p)
        elif v == 3:
            ngrid.add(p)
    return ngrid