Advent of code 2020, Day 17, Python3.9

day17.py

from sys import stdin
from itertools import product
from collections import defaultdict


def build(lines, dims):
    return {(x, y)+(0,)*(dims-2)
            for x, l in enumerate(lines)
            for y, v in enumerate(l.strip())
            if v == '#'}


def play(world, dims):
    sums = defaultdict(int)
    for w in (t for x in world
              for t in product(*[[x[i]-1, x[i], x[i]+1]
                                 for i in range(dims)])
              if any(t[j] != x[j] for j in range(dims))):
        sums[w] += 1
    return {w for w in {w for x in world
                        for w in product(*[[x[i]-1, x[i], x[i]+1]
                                           for i in range(dims)])}
            if sums[w] == 3 and w not in world or (
        w in world and (sums[w] in (2, 3)))}


I = [l for l in stdin]
for dims in (3, 4):
    world = build(I, dims)
    for _ in range(6):
        world = play(world, dims)
    print(len(world))