Advent of code 2023 day 22

def day22(s, *, part2=False):
  bricks = [tuple(map(int, re.findall(r'\d+', line))) for line in s.splitlines()]
  bricks.sort(key=lambda brick: brick[2])  # In ascending minimum height.

  heights = np.full((10, 10), 0)
  occupancy = np.full((10, 10, max(brick[5] for brick in bricks) + 1), -1)  # Empty.
  occupancy[..., 0] = -2  # Fake extra base brick spanning height=0.
  placement, belows = {}, {}
  total = 0

  for index, brick in enumerate(bricks):
    footprint = slice(brick[0], brick[3] + 1), slice(brick[1], brick[4] + 1)
    height = heights[footprint].max() + 1
    placement[index] = height
    occupancy[(*footprint, slice(height, height + 1 + brick[5] - brick[2]))] = index
    belows[index] = set(occupancy[(*footprint, height - 1)].flat) - {-1}
    heights[footprint] = height + brick[5] - brick[2]

  if not part2:
    for index, brick in enumerate(bricks):
      footprint = slice(brick[0], brick[3] + 1), slice(brick[1], brick[4] + 1)
      height = placement[index]
      height2 = height + 1 + brick[5] - brick[2]
      # Check if for each brick  above this one, there is at least one other brick below it.
      above = set(occupancy[(*footprint, height2)].flat) - {-1}
      total += all(len(belows[index2]) > 1 for index2 in above)
    return total

  for index in range(len(bricks)):
    falling = {index}
    for index2 in range(index + 1, len(bricks)):
      if belows[index2] <= falling:
        falling.add(index2)
    total += len(falling) - 1
  return total