2.py (2082B) [raw]
1 #!/usr/bin/env python3 2 import sys 3 4 def expand(dense): 5 fs = [] 6 for i in range(len(dense)): 7 if i % 2 == 0: 8 for j in range(int(dense[i])): 9 fs.append(i // 2) 10 else: 11 for j in range(int(dense[i])): 12 fs.append(None) 13 return fs 14 15 def pprint(fs): 16 for c in fs: 17 if c is not None: 18 print(c, end="") 19 else: 20 print(".", end="") 21 print("") 22 23 def compact(fs): 24 comp = [c for c in fs] 25 26 # First store the free space to avoid rescanning every time 27 free = {} 28 start = None 29 size = 0 30 for i in range(len(comp)): 31 if comp[i] is not None: 32 if start is not None: 33 free[start] = size 34 start = None 35 size = 0 36 else: 37 if start is None: 38 start = i 39 size += 1 40 41 # work backwards 42 id = comp[-1] 43 end = len(comp) - 1 44 for i in range(len(comp) - 1, -1, -1): 45 if comp[i] == id: 46 # still in file block 47 continue 48 else: 49 if id is not None: 50 # span from i + 1 -> end needs moving 51 start = i + 1 52 size = end - start + 1 53 for f in sorted(free.keys()): 54 if f < start and free[f] >= size: 55 # copy the block 56 for j in range(size): 57 comp[f + j] = comp[start + j] 58 comp[start + j] = None 59 # update our space 60 free[f + size] = free[f] - size 61 del(free[f]) 62 break 63 # reset state 64 id = comp[i] 65 end = i 66 return comp 67 68 def chksum(fs): 69 tot = 0 70 for i in range(len(fs)): 71 if fs[i] is None: 72 continue 73 tot += fs[i] * i 74 75 return tot 76 77 if __name__ == '__main__': 78 dense = next(sys.stdin).strip() 79 fs = expand(dense) 80 comp = compact(fs) 81 # pprint(comp) 82 print(chksum(comp)) 83