aoc2024/d09/part2.py

31 lines
630 B
Python
Raw Normal View History

import sys, os
I = list(map(int, sys.stdin.read().strip()))
def sum_n(a, b):
return (b-a+1)*(a+b)//2
idxs = []
b = 0
for i in I:
idxs.append(b)
b += i
res = 0
id = len(I)//2 # last id
for j in range(len(I)-1,-1,-2): # backward
fsize = I[j]
for i in range(1, j, 2): # find free space forward
if I[i] >= fsize: # enough free space
b = idxs[i]
res += sum_n(b, b+fsize-1)*id
idxs[i] += fsize
I[i] -= fsize
break
else: # block too big stays where it is
b = idxs[j]
res += sum_n(b, b+fsize-1)*id
id -= 1
print(res)