day17.py

import numpy as np
from scipy.ndimage import convolve

LINES = open("inputs/day17").read().splitlines()

dims = 3  # 4 for part 2
steps = 6
front = 2 * steps + 1
filt = np.full((3,) * dims, 1)
filt[(1,) * dims] = 100

board = np.zeros((front,) * (dims - 2) + (front + len(LINES), front + len(LINES[0])), dtype=np.uint8)
for i, l in enumerate(LINES):
    for j, c in enumerate(l):
        board[(steps,) * (dims - 2) + (steps + i, steps + j)] = c == "#"

for i in range(steps):
    conv = convolve(board, filt, mode="constant")
    board[:] = (conv == 3) | (conv == 102) | (conv == 103)

print(np.sum(board))