https://adventofcode.com/2023/day/22

day22.py

import itertools as it
import re

cubes = sorted(([*map(int, re.findall(r'\d+', l))] for l in open(0)),
               key=lambda p: p[2])


def points(c, dz=0):
    x1, y1, z1, x2, y2, z2 = c
    return it.product(range(x1, x2+1), range(y1, y2+1), range(z1-dz, z2+1-dz))


def fall(cs, fall_from=0, occupied=set()):
    fallen, moved_cnt = [], 0
    occupied = occupied.copy()
    for c in cs[fall_from:]:
        x1, y1, z1, x2, y2, z2 = c
        assert x1<=x2 and y1<=y2 and z1<=z2
        dz = 0
        while z1-dz>1 and occupied.isdisjoint(points(c, dz+1)):
            dz += 1
        fallen += [(x1, y1, z1-dz, x2, y2, z2-dz)]
        occupied.update(points(c, dz))
        moved_cnt += dz>0
    return sorted(fallen, key=lambda p: p[2]), moved_cnt, occupied


settled, _, occupied = fall(cubes)
ans = 0
ans2 = 0
for i in reversed(range(len(settled))):
    occupied.difference_update(points(settled[i]))
    _, moved_cnt, _ = fall(settled[:i]+settled[i+1:], i, occupied)
    ans += moved_cnt==0
    ans2 += moved_cnt

print(ans)
print(ans2)