85 lines
2 KiB
Python
85 lines
2 KiB
Python
from pathlib import Path
|
|
import numpy as np
|
|
|
|
|
|
class HardDrive:
|
|
def __init__(self, input, verbose=False):
|
|
self.volume = self._arrange_data(input)
|
|
self.du = sum(~np.isnan(self.volume))
|
|
self.verbose = verbose
|
|
if verbose:
|
|
print(input)
|
|
print(self)
|
|
|
|
def __str__(self):
|
|
a = ""
|
|
for i in self.volume:
|
|
if np.isnan(i):
|
|
a += "."
|
|
else:
|
|
a += str(int(i))
|
|
return a
|
|
|
|
@staticmethod
|
|
def _arrange_data(input):
|
|
lens = np.array([int(i) for i in input])
|
|
len_mem = np.sum(lens)
|
|
vol = np.zeros(len_mem)
|
|
|
|
i = 0
|
|
s = 1
|
|
d = 0
|
|
for j in lens:
|
|
if s > 0:
|
|
vol[i : i + j] = d
|
|
d += 1
|
|
else:
|
|
vol[i : i + j] = np.nan
|
|
s *= -1
|
|
i += j
|
|
return vol
|
|
|
|
@property
|
|
def reverse(self):
|
|
a = self.volume[~np.isnan(self.volume)]
|
|
return a[::-1]
|
|
|
|
def frag_sort(self):
|
|
a = self.volume[:]
|
|
r = self.reverse[:]
|
|
t = np.isnan(self.volume)
|
|
a[t] = r[: sum(t)]
|
|
a[self.du :] = np.nan
|
|
self.volume = a
|
|
if self.verbose:
|
|
print(self)
|
|
|
|
@property
|
|
def holes(self):
|
|
(i,) = np.where(np.isnan(self.volume))
|
|
d = np.diff(i, prepend=0) != 1
|
|
print(np.diff(i, prepend=0) == 1)
|
|
# TODO get hole width!
|
|
return i[d]
|
|
|
|
def clean_sort(self):
|
|
self.holes
|
|
|
|
@property
|
|
def checksum(self):
|
|
a = np.array(range(len(self.volume)))
|
|
i = ~np.isnan(self.volume)
|
|
return np.sum(self.volume[i] * a[i], dtype=int)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
filepath = Path("./example")
|
|
# filepath = Path("./input")
|
|
|
|
with open(filepath, "r") as filein:
|
|
data = filein.read().rstrip()
|
|
disk = HardDrive(data, verbose=True)
|
|
print(disk.checksum)
|
|
# disk.frag_sort()
|
|
disk.clean_sort()
|
|
print(disk.checksum)
|