aoc

Advent of Code Solutions
git clone git://git.alexkarle.com.com/aoc
Log | Files | Refs | README | LICENSE

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