day 9 part 1

day09/sol.py

@@ -0,0 +1,45 @@
+sample = "2333133121414131402"
+input = open("input").read().strip()
+
+def expand(s):
+    disk = []
+    id = 0
+    for i in range(len(s)):
+        d = int(s[i])
+        if i % 2 == 0:
+            disk.extend([id]*d)
+            id += 1
+        else:
+            disk.extend([-1] * d) # free space
+    return disk
+
+def solve(input):
+    disk = expand(input)
+    if len(disk) < 100:
+        print(disk)
+    while disk[-1] == -1:
+        disk.pop()
+    free = 0
+    last = len(disk)-1
+    while free < last:
+        while disk[free] != -1:
+            free += 1
+        if free >= last:
+            break
+        disk[free] = disk[last]
+        disk[last] = -1
+        while free < last and disk[last] == -1:
+            last -= 1
+
+        if len(disk) < 100:
+            print(free, last, disk)
+    t = 0
+    for i, x in enumerate(disk):
+        if x == -1:
+            continue
+        t += i*x
+    print(t)
+
+
+solve(sample)
+solve(input)