AOC2021 Day 22 (Python)

day22_tests.py

import re


line_format = re.compile(r'(on|off) x=(-?\d+)..(-?\d+),y=(-?\d+)..(-?\d+),z=(-?\d+)..(-?\d+)')


def read_line(line):
    instr, *bounds = line_format.search(line).groups()
    return instr, tuple(int(bound) for bound in bounds)


def clip(box_a, box_b):
    min_x_a, max_x_a, min_y_a, max_y_a, min_z_a, max_z_a = box_a
    min_x_b, max_x_b, min_y_b, max_y_b, min_z_b, max_z_b = box_b
    
    clipped_min_x = max(min_x_a, min_x_b)
    clipped_max_x = min(max_x_a, max_x_b)
    clipped_min_y = max(min_y_a, min_y_b)
    clipped_max_y = min(max_y_a, max_y_b)
    clipped_min_z = max(min_z_a, min_z_b)
    clipped_max_z = min(max_z_a, max_z_b)

    if clipped_min_x > clipped_max_x or clipped_min_y > clipped_max_y or clipped_min_z > clipped_max_z:
        return None

    clipped = clipped_min_x, clipped_max_x, clipped_min_y, clipped_max_y, clipped_min_z, clipped_max_z
    return clipped


def clip_all(clip_box, boxen):
    return list(set(clipped for clipped in (clip(box, clip_box) for box in boxen) if clipped is not None))


def box_volume(box):
    min_x, max_x, min_y, max_y, min_z, max_z = box
    return (max_x + 1 - min_x) * (max_y + 1 - min_y) * (max_z + 1 - min_z)


def sum_volume(boxen):
    if len(boxen) == 0:
        return 0
    first, *remainder = boxen
    overlaps = clip_all(first, remainder)
    return box_volume(first) + sum_volume(remainder) - sum_volume(overlaps)


def count_lit_cubes(typed_boxen):
    if len(typed_boxen) == 0:
        return 0

    (box_type, first), *remainder = typed_boxen
    if box_type == 'off':
        return count_lit_cubes(remainder)

    overlaps = clip_all(
        first,
        (box for _, box in remainder))

    return box_volume(first) + count_lit_cubes(remainder) - sum_volume(overlaps)


def read_file(file):
    return [read_line(line.strip()) for line in file]


def test_within_bounds():

    example1 = """
on x=10..12,y=10..12,z=10..12
on x=11..13,y=11..13,z=11..13
off x=9..11,y=9..11,z=9..11
on x=10..10,y=10..10,z=10..10
    """.strip().split('\n')

    assert count_lit_cubes(read_file(example1)) == 39

    with open('puzzle_inputs/day22.txt') as file:
        boxen = read_file(file)

    clipping_box = (-50, 50, -50, 50, -50, 50)
    clipped = [(box_type, clipped_box)
               for box_type, clipped_box in ((box_type, clip(clipping_box, box))
                                             for box_type, box in boxen)
               if clipped_box is not None]
    assert count_lit_cubes(clipped) == 607657

    print()
    print(count_lit_cubes(boxen))