Implement 2024 day 9

Author: bertptrs

2024/src/aoc/days/day9.py

@@ -0,0 +1,107 @@
+import heapq
+
+from . import SeparateRunner
+
+
+def file_checksum(file_id: int, start: int, length: int) -> int:
+    return file_id * length * (2 * start + length - 1) // 2
+
+
+class DayRunner(SeparateRunner):
+    @classmethod
+    def part1(cls, input: str) -> int:
+        files = []
+        empty = []
+
+        pos = 0
+
+        for c in input.strip():
+            val = int(c)
+
+            if len(files) == len(empty):
+                files.append((pos, val))
+            else:
+                empty.append((pos, val))
+
+            pos += val
+
+        checksum = 0
+
+        for start, length in empty:
+            while files and length > 0:
+                file_start, file_len = files.pop()
+                if file_start < start:
+                    files.append((file_start, file_len))
+                    break
+
+                file_id = len(files)
+
+                infill = min(file_len, length)
+
+                checksum += file_checksum(file_id, start, infill)
+                start += infill
+
+                if infill != file_len:
+                    files.append((file_start, file_len - infill))
+
+                length -= infill
+            else:
+                continue
+            break
+
+        for file_id, (file_start, file_len) in enumerate(files):
+            checksum += file_checksum(file_id, file_start, file_len)
+
+        return checksum
+
+    @classmethod
+    def part2(cls, input: str) -> int:
+        files = []
+        empty = [[] for _ in range(10)]
+
+        pos = 0
+
+        is_file = True
+
+        for c in input.strip():
+            val = int(c)
+
+            if is_file:
+                files.append((pos, val))
+                is_file = False
+            else:
+                # No need for heappush, as we're appending values in order
+                empty[val].append(pos)
+                is_file = True
+
+            pos += val
+
+        checksum = 0
+
+        while files:
+            start, length = files.pop()
+            file_id = len(files)
+
+            best = None
+            best_heap = None
+
+            for i, heap in enumerate(empty[length:]):
+                if not heap or heap[0] > start:
+                    continue
+
+                if best is None or best > heap[0]:
+                    best = heap[0]
+                    best_heap = i + length
+
+            if best is None:
+                # No room to move left, count score at current position
+                checksum += file_checksum(file_id, start, length)
+            else:
+                checksum += file_checksum(file_id, best, length)
+                heapq.heappop(empty[best_heap])
+
+                if length < best_heap:
+                    remainder = best_heap - length
+                    heapq.heappush(empty[remainder], best + length)
+
+        return checksum

2024/tests/test_day09.py

@@ -0,0 +1,11 @@
+from aoc.days.day9 import DayRunner
+
+SAMPLE = "2333133121414131402"
+
+
+def test_sample_part1() -> None:
+    assert DayRunner.part1(SAMPLE) == 1928
+
+
+def test_sample_part2() -> None:
+    assert DayRunner.part2(SAMPLE) == 2858