diff --git a/2024/day9/day9.py b/2024/day9/day9.py new file mode 100644 index 0000000..6753fb0 --- /dev/null +++ b/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) +