AoC 2023 day 22

import numpy as np
from re import findall
from collections import deque


def find_all_ints(line):
    return list(map(int, findall(r'\d+', line)))


def read_input():
    with open('../Inputs/day_22_input.txt') as f:
        blocks = [find_all_ints(line) for line in f.read().strip().splitlines()]
        blocks.sort(key=lambda u: u[2])
        # Assume lower bounds are given before upper bounds.
        max_x = max(block[3] for block in blocks)
        max_y = max(block[4] for block in blocks)
        max_z = max(block[5] for block in blocks)
        block_places = np.zeros((max_x + 1, max_y + 1, max_z + 1), dtype=bool)
        for block in blocks:
            block_places[block[0]:block[3] + 1, block[1]:block[4] + 1, block[2]:block[5] + 1] = True
        block_places[:, :, 0] = True
    return blocks, block_places


def drop_blocks(blocks, block_places):
    for block in blocks:
        x1, y1, z1, x2, y2, z2 = block
        section = block_places[x1:x2 + 1, y1:y2 + 1, :z1]
        flattened_section = section.sum(axis=0).sum(axis=0)
        support_height = np.nonzero(flattened_section)[0].max()
        fall_height = z1 - (support_height + 1)
        block_places[x1:x2 + 1, y1:y2 + 1, z1:z2 + 1] = False
        block_places[x1:x2 + 1, y1:y2 + 1, z1 - fall_height:z2 - fall_height + 1] = True
        block[2] -= fall_height
        block[5] -= fall_height
    blocks.sort(key=lambda u: u[2])
    return blocks, block_places


def get_blocks_by_z(blocks):
    max_z = max(block[5] for block in blocks)
    # By making these one longer than necessary, we can
    # always safely ask for the element with index block[5] + 1.
    blocks_by_lowest_z = [[] for _ in range(max_z + 2)]
    blocks_by_highest_z = [[] for _ in range(max_z + 2)]
    for block in blocks:
        blocks_by_lowest_z[block[2]].append(block)
        blocks_by_highest_z[block[5]].append(block)
    return blocks_by_lowest_z, blocks_by_highest_z


def scan_below(blocks, blocks_by_highest_z, block_indices):
    # Save the blocks as their index, since this will reduce hash overhead compared to
    # tuple or string of the six numbers that represent a block.
    # Since blocks is sorted by the lowest z coordinate of the blocks, is_supported_by will have the same ordering.
    is_supported_by = [set() for i, _ in enumerate(blocks)]
    for i, block_1 in enumerate(blocks):
        x1, y1, z1, x2, y2, z2 = block_1
        # The lowest z coordinate of block_1 needs to be 1 greater than the highest z coordinate of block_2
        # to have the latter support the former.
        for block_2 in blocks_by_highest_z[z1 - 1]:
            # Closed intervals [a, b] and [c, d] overlap if and only if a <= d and c <= b.
            if x1 <= block_2[3] and block_2[0] <= x2 and y1 <= block_2[4] and block_2[1] <= y2:
                is_supported_by[i].add(block_indices[tuple(block_2)])
    return is_supported_by


def scan_above(blocks, blocks_by_lowest_z, block_indices):
    # Save the blocks as their index, since this will reduce hash overhead compared to
    # tuple or string of the six numbers that represent a block.
    # Since blocks is sorted by the lowest z coordinate of the blocks, is_supporting will have the same ordering.
    is_supporting = [set() for i, _ in enumerate(blocks)]
    for i, block_1 in enumerate(blocks):
        x1, y1, z1, x2, y2, z2 = block_1
        # The lowest z coordinate of block_2 needs to be 1 greater than the highest z coordinate of block_1
        # to have the latter support the former.
        for block_2 in blocks_by_lowest_z[z2 + 1]:
            # Closed intervals [a, b] and [c, d] overlap if and only if a <= d and c <= b.
            if x1 <= block_2[3] and block_2[0] <= x2 and y1 <= block_2[4] and block_2[1] <= y2:
                is_supporting[i].add(block_indices[tuple(block_2)])
    return is_supporting


def find_destroyable_blocks(is_supported_by):
    # Blocks can be safely destroyed if they are not the sole support of some other block.
    not_destroyable = set()
    for supports in is_supported_by:
        if len(supports) == 1:
            support, = supports
            not_destroyable.add(support)
    # is_supported_by has an element for every block, so this is the number of destroyable blocks.
    return len(is_supported_by) - len(not_destroyable)


def find_falling_num(block, is_supporting, is_supported_by):
    # A block does not really support itself but this makes the checks cleaner.
    indirectly_supporting = {block}
    # is_supporting is sorted by lowest z coordinate. This means the blocks added onto the queue
    # are ordered as well and that gives the correct order for the support check.
    queue = deque(is_supporting[block])
    while queue:
        on_top = queue.popleft()
        if is_supported_by[on_top].issubset(indirectly_supporting):
            indirectly_supporting.add(on_top)
        queue.extend(is_supporting[on_top])
    # Take 1 away to counteract the block itself being added.
    return len(indirectly_supporting) - 1


def main():
    blocks, block_places = read_input()
    blocks, block_places = drop_blocks(blocks, block_places)
    block_indices = {tuple(block): i for i, block in enumerate(blocks)}
    blocks_by_lowest_z, block_by_highest_z = get_blocks_by_z(blocks)
    is_supported_by = scan_below(blocks, block_by_highest_z, block_indices)
    destroyable = find_destroyable_blocks(is_supported_by)
    print(f'Part one: {destroyable}')
    is_supporting = scan_above(blocks, blocks_by_lowest_z, block_indices)
    total_falling = sum(find_falling_num(i, is_supporting, is_supported_by) for i, _ in enumerate(blocks))
    print(f'Part two: {total_falling}')


if __name__ == '__main__':
    from time import perf_counter as pf
    t0 = pf()
    main()
    print(pf() - t0)