AoC 2021 Day22 part2

AoC 2021 Day22 part2.py

from itertools import product

cmd_blocks = []  # list of (set, cube)
for li in open('day22.txt'):
    # on x=-11..33,y=-6..40,z=-16..37
    cmd, rblocks = li.split()
    xs, ys, zs = rblocks.split(',')
    # print(xs, ys, zs)
    xs = [int(x) for x in xs.split('x=')[1].split('..')]
    ys = [int(x) for x in ys.split('y=')[1].split('..')]
    zs = [int(x) for x in zs.split('z=')[1].split('..')]
    # make range end exclusive
    xs[1] += 1
    ys[1] += 1
    zs[1] += 1
    cmd_blocks.append((cmd == 'on', (zs, ys, xs)))
    
    
def subtract_cube(on_cube, off_cube):
    # if off cube overlaps with on_cube, split on_cube into 26(!) cubes to 
    # account for the missing cube in the middle (general case)
    overlaps = []
    for onx, offx in zip(on_cube, off_cube):
        overlap_r = max(onx[0], offx[0]), min(onx[1], offx[1])
        if overlap_r[0] >= overlap_r[1]:
            yield on_cube # no overlap
            return
        overlaps.append(overlap_r)
    # for each side we have cube_start, overlap_start, overlap_end, cube_end
    ranges = [
        [(cube[0], overlap[0]),
        (overlap[0], overlap[1]),
        (overlap[1], cube[1])] for cube, overlap in zip(on_cube, overlaps)
               ]
    # we have 3^N sub cubes. Want to skip the middle one
    mid_point = (3 ** len(ranges)) // 2
    for i, rs in enumerate(product(*ranges)):
        if i == mid_point:
            continue
        if sum_cube(rs) != 0:
            yield rs
                    

def sum_cube(c):
    r = 1
    for pair in c:
        r = r * (pair[1] - pair[0])
    return r

on_cubes = []
for turn_on, newc in cmd_blocks:
    # remove from existing:
    new_on_cubes = [c for onc in on_cubes
                    for c in subtract_cube(onc, newc)]
    if turn_on:
        new_on_cubes.append(newc)
    on_cubes = new_on_cubes
    
print(sum(sum_cube(c) for c in on_cubes))