day 9 part 2

it seemed easier to go for a sparse implementation than try to make the
dense implementation from part 1 work

day09/sol.py

@@ -0,0 +1,85 @@
+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)
+
+def solve2(input):
+    freelist = []
+    blocks = {}
+    pos = 0
+    id = 0
+    for i in range(len(input)):
+        n = int(input[i])
+        if i % 2 == 0:
+            blocks[id] = (pos, n)
+            id += 1
+        else:
+            freelist.append((pos, n))
+        pos += n
+
+    lastid = id
+    for id in range(lastid-1,-1,-1):
+        bpos, blen = blocks[id]
+        for i in range(len(freelist)):
+            fpos, flen = freelist[i]
+            if bpos < fpos:
+                break
+            if blen <= flen:
+                blocks[id] = (fpos, blen)
+                if flen == blen:
+                    del freelist[i]
+                else:
+                    freelist[i] = (fpos+blen, flen-blen)
+                break
+
+    t = 0
+    for id in blocks:
+        i, n = blocks[id]
+        for x in range(i,i+n):
+            t += id*x
+
+    print(t)
+
+
+solve(sample)
+solve(input)
+
+solve2(sample)
+solve2(input)