2024: d09: ex2: add solution

2024/d09/ex2/ex2.py

@@ -0,0 +1,93 @@
+#!/usr/bin/env python
+
+import sys
+from typing import NamedTuple
+
+
+class FilePosition(NamedTuple):
+    pos: int
+    length: int
+
+
+class ParsedDisk(NamedTuple):
+    files: list[FilePosition]
+    holes: list[FilePosition]
+
+
+def solve(input: str) -> int:
+    def parse(input: str) -> tuple[list[int], list[int]]:
+        files: list[int] = []
+        free: list[int] = []
+
+        for i, c in enumerate(input):
+            if c == "\n":
+                continue
+            (files if i % 2 == 0 else free).append(int(c))
+
+        # Make `to_disk` slightly simpler
+        if len(files) > len(free):
+            free.append(0)
+
+        return files, free
+
+    def to_disk(files: list[int], free: list[int]) -> ParsedDisk:
+        assert len(files) == len(free)  # Sanity check
+
+        pos = 0
+        disk_files: list[FilePosition] = []
+        holes: list[FilePosition] = []
+
+        for i in range(len(files)):
+            disk_files.append(FilePosition(pos, files[i]))
+            pos += files[i]
+            holes.append(FilePosition(pos, free[i]))
+            pos += free[i]
+
+        return ParsedDisk(disk_files, holes)
+
+    def compact(disk: ParsedDisk) -> list[int | None]:
+        def move_files(disk: ParsedDisk) -> dict[int, FilePosition]:
+            new_files: dict[int, FilePosition] = {}
+            for i, (pos, length) in reversed(list(enumerate(disk.files))):
+                for h in range(len(disk.holes)):
+                    # Hole must be big enough to fit the file
+                    if length > disk.holes[h].length:
+                        continue
+                    # Hole must be to the left of the file
+                    if pos < disk.holes[h].pos:
+                        break
+                    # We found a hole, move the file into it
+                    pos = disk.holes[h].pos
+                    disk.holes[h] = FilePosition(
+                        length=disk.holes[h].length - length,
+                        pos=disk.holes[h].pos + length,
+                    )
+                    break
+                new_files[i] = FilePosition(pos, length)
+            return new_files
+
+        def to_disk(files: dict[int, FilePosition]) -> list[int | None]:
+            disk: list[int | None] = [None] * max(
+                (f.pos + f.length) for f in files.values()
+            )
+            for i, f in files.items():
+                for j in range(f.length):
+                    assert disk[f.pos + j] is None  # Sanity check
+                    disk[f.pos + j] = i
+            return disk
+
+        new_files = move_files(disk)
+        return to_disk(new_files)
+
+    files, free = parse(input)
+    disk = to_disk(files, free)
+    return sum(i * n for i, n in enumerate(compact(disk)) if n is not None)
+
+
+def main() -> None:
+    input = sys.stdin.read()
+    print(solve(input))
+
+
+if __name__ == "__main__":
+    main()