Day 22 - shortened code, more precise annotations

day22.py

@@ -0,0 +1,74 @@
+from collections import defaultdict
+from tools.aoc import AOCDay
+from typing import Any
+
+
+def get_space_below(brick: tuple[tuple[int, int, int], ...], max_y: dict[tuple[int, int], int]) -> int:
+    below_y = 0
+    for x in range(brick[0][0], brick[-1][0] + 1):
+        for z in range(brick[0][2], brick[-1][2] + 1):
+            below_y = max(below_y, max_y[(x, z)])
+
+    return brick[0][1] - below_y - 1
+
+
+def drop(bricks: set[tuple[tuple[int, int, int], ...]]) -> (set[tuple[tuple[int, int, int], ...]], int):
+    dropped_bricks = set()
+    max_y = defaultdict(int)
+    count = 0
+    for brick in sorted(bricks, key=lambda b: b[0][1]):
+        space_below = get_space_below(brick, max_y)
+        new_brick = tuple()
+        for c in brick:
+            new_y = c[1] - space_below
+            max_y[(c[0], c[2])] = max(max_y[(c[0], c[2])], new_y)
+            new_brick += ((c[0], new_y, c[2]),)
+
+        dropped_bricks.add(new_brick)
+        if space_below:
+            count += 1
+
+    return dropped_bricks, count
+
+
+class Day(AOCDay):
+    inputs = [
+        [
+            (5, "input22_test"),
+            (527, "input22"),
+        ],
+        [
+            (7, "input22_test"),
+            (100376, "input22"),
+        ]
+    ]
+
+    def parse_input(self) -> set[tuple[tuple[int, int, int], ...]]:
+        bricks = set()
+        for line in self.getInput():
+            a, b = line.split("~")
+            x1, z1, y1 = map(int, a.split(","))
+            x2, z2, y2 = map(int, b.split(","))
+            bricks.add(
+                tuple(
+                    (x, y, z)
+                    for x in range(x1, x2 + 1)
+                    for y in range(y1, y2 + 1)
+                    for z in range(z1, z2 + 1)
+                )
+            )
+
+        return bricks
+
+    def part1(self) -> Any:
+        dropped_bricks, _ = drop(self.parse_input())
+        return sum(c[1] == 0 for c in (drop(dropped_bricks - {x}) for x in dropped_bricks))
+
+    def part2(self) -> Any:
+        dropped_bricks, _ = drop(self.parse_input())
+        return sum(c[1] for c in (drop(dropped_bricks - {x}) for x in dropped_bricks))
+
+
+if __name__ == '__main__':
+    day = Day(2023, 22)
+    day.run(verbose=True)

inputs/input22_test

@@ -0,0 +1,7 @@
+1,0,1~1,2,1
+0,0,2~2,0,2
+0,2,3~2,2,3
+0,0,4~0,2,4
+2,0,5~2,2,5
+0,1,6~2,1,6
+1,1,8~1,1,9