[wip] 2024 day 9

2024/day9/day9.py

@@ -0,0 +1,103 @@
+from dataclasses import dataclass
+from typing import Literal
+
+
+@dataclass()
+class Item:
+    kind: Literal["file", "free"]
+    size: int
+    id_: int = 0
+
+
+def parse_disk(data):
+    disk = {}
+    disk2 = []
+    id_cnt = 0
+    offset = 0
+    for pos, size in enumerate(data):
+        if pos % 2 == 0:  # file
+            id = id_cnt
+            for i in range(int(size)):
+                disk[offset + i] = id
+            disk2.append(Item("file", int(size), id))
+            id_cnt += 1
+            offset += int(size)
+        else:  # free space
+            if int(size) > 0:
+                disk2.append(Item("free", int(size)))
+                offset += int(size)
+    return disk, disk2
+
+
+def main(inp):
+    disk, disk2 = parse_disk(inp)
+    max_file_location = max(disk.keys())
+
+    write_ptr = 0
+    read_ptr = max_file_location
+    while True:
+        while write_ptr in disk:
+            write_ptr += 1
+
+        while read_ptr not in disk:
+            read_ptr -= 1
+        
+        if write_ptr >= read_ptr:
+            break
+
+        disk[write_ptr] = disk[read_ptr]
+        del disk[read_ptr]
+
+    checksum = sum(i * disk.get(i, 0) for i in range(max_file_location))
+    print("Part 1: ", checksum)
+
+    print(len(disk2))
+    write_ptr = 0
+    read_ptr = len(disk2) - 1
+    for file_index in range(len(disk2) - 1, 0, -1):
+        file = disk2[file_index]
+        if file.kind != "file":
+            continue
+
+        # find index of the first gap large enough
+        free_index, free_space = next(((i, b) for i, b in enumerate(disk2) if b.kind == "free" and b.size >= file.size), (None, None))
+        if free_index is None:
+            continue
+        # add a free space in place of the file
+        disk2[file_index] = Item("free", file.size)
+        # insert file just before free space
+        disk2.insert(free_index, file)
+        # decrease free space by file size (in case free space is larger)
+        free_space.size -= file.size
+    total_checksum = 0
+    offset = 0
+    print(disk2)
+    print(len(disk2))
+    for f in disk2:
+        if f.kind != "file":
+            offset += f.size
+            continue
+        # S(n) = n*(n+1) // 2
+        print(f"checksum = {f.id_} * ({offset} * {f.size} + ({f.size} * ({f.size - 1})) // 2")
+        checksum = f.id_ * (offset * f.size + (f.size * (f.size - 1)) // 2)
+        print(f, checksum)
+
+        offset += f.size
+        total_checksum += checksum
+    print(total_checksum)
+
+
+
+        
+
+
+
+if __name__ == "__main__":
+    import sys
+    infile = sys.argv[1] if len(sys.argv) > 1 else None
+    content = "2333133121414131402"
+    if infile is not None:
+        with open(infile) as f:
+            content = f.read().rstrip()
+    main(content)
+