17.2.py

from itertools import product
from operator import itemgetter
with open('input/17.txt') as f:
  data = [[y for y in x.strip()] for x in f.readlines()]

def get_living_neighbours(cube, point):
  # get neighbour ranges
  ranges = []
  for p in point:
    ranges.append((p - 1, p, p + 1))

  # count living cells
  alive = 0
  for p in product(*ranges):
    if p in cube:
      alive += 1

  # remove self
  if point in cube:
    alive -= 1

  return alive

def extend(cube):
  # for each dimension extend the lowest
  # and highest value by one
  ranges = []
  for i in range(4):
    maximum = max(map(itemgetter(i), cube)) + 1
    minimum = min(map(itemgetter(i), cube)) - 1
    ranges.append(range(minimum, maximum+1))
  return product(*ranges)

def cycle(cube):
  new_cube = set()
  # get the valid range for an extended cube in each dimension
  # use that to form our new cube
  for point in extend(cube):
    alive = get_living_neighbours(cube, point)
    if (point in cube and alive == 2) or alive == 3:
      new_cube.add(point)
  return new_cube

width = len(data)
# represent our cube as only the living points
# now we have 4 dimensions so fill 2 of them with 0
cube = set()
for x, y in product(range(width), range(width)):
  if data[x][y] == '#':
    cube.add((0, 0, x, y))

for i in range(6):
  cube = cycle(cube)
print("Part 2:", len(cube))