Day 09

day09.py

@@ -0,0 +1,125 @@
+#!/usr/bin/env python3
+from copy import copy
+from typing import Optional
+from dataclasses import dataclass
+
+
+@dataclass
+class Block:
+    fid: Optional[int]
+    length: int
+
+
+puzzle = []
+with open("day09") as f:
+    puzzle = list(map(int, f.readline().strip()))
+
+fs = []
+
+def fill_fs():
+    fid = 0
+    for i,e in enumerate(puzzle):
+        if i%2==0: # file
+            fs.append(Block(fid, e))
+            fid += 1
+        else: # gap
+            fs.append(Block(None, e))
+
+def find_gaps():
+    gaps = []
+    for i,b in enumerate(fs):
+        if b.fid is None: gaps.append((i, b))
+    return gaps
+
+def find_files():
+    files = []
+    for i,b in enumerate(fs):
+        if b.fid is not None: files.append((i, b))
+    return files
+
+def checksum():
+    checksum = 0
+    idx = 0
+    for b in fs:
+        if b.fid is None:
+            for i in range(b.length): idx += 1
+        else:
+            for i in range(b.length):
+                checksum += b.fid*idx
+                idx += 1
+
+    return checksum
+
+def part1():
+    while True:
+        gaps = find_gaps()
+        if not gaps: break
+
+        if gaps[-1][0] == len(fs)-1: # free block at end
+            del fs[-1]
+            continue
+
+        files = find_files()
+        next_gap = gaps[0][0]
+        next_file = files[-1][0]
+
+        if fs[next_gap].length == fs[next_file].length:
+            fs[next_gap] = fs[next_file]
+            del fs[-1]
+            continue
+
+        if fs[next_gap].length > fs[next_file].length:
+            gap_length = fs[next_gap].length
+            file_length = fs[next_file].length
+
+            fs.insert(next_gap, fs[next_file])
+            fs[next_gap+1].length = gap_length-file_length
+            del fs[-1]
+
+            continue
+
+        if fs[next_gap].length < fs[next_file].length:
+            gap_length = fs[next_gap].length
+            file_length = fs[next_file].length
+            fid = fs[next_file].fid
+
+            fs[next_file].length = file_length - gap_length
+            fs[next_gap] = Block(fid=fid, length=gap_length)
+
+            continue
+
+fill_fs()
+part1()
+print(checksum())
+
+def part2():
+    files = list(reversed(find_files()))
+
+    for i,file in enumerate(files):
+        gaps = find_gaps()
+        if not gaps: break
+
+        file_idx = fs.index(file[1])
+
+        for gap in gaps:
+            if gap[0] >= file_idx:
+                break
+
+            if gap[1].length == file[1].length:
+                fs[file_idx] = Block(fid=None, length=file[1].length)
+                fs[gap[0]] = file[1]
+                break
+
+            if gap[1].length > file[1].length:
+                gap_length = gap[1].length
+                file_length = file[1].length
+
+                fs[file_idx] = Block(fid=None, length=file[1].length)
+
+                fs.insert(gap[0], file[1])
+                fs[gap[0]+1].length = gap_length-file_length
+                break
+
+fill_fs()
+part2()
+print(checksum())