129 lines
No EOL
3.4 KiB
Python
129 lines
No EOL
3.4 KiB
Python
import sys
|
|
from collections import namedtuple, defaultdict
|
|
|
|
|
|
Segment = namedtuple("Segment", ["id", "len"])
|
|
|
|
|
|
def parse(instr: str) -> list[Segment]:
|
|
res = []
|
|
next_id = 0
|
|
|
|
for i, v in enumerate(instr):
|
|
v = int(v)
|
|
if i % 2 == 0:
|
|
# this is a file
|
|
res.append(Segment(next_id, v))
|
|
next_id += 1
|
|
else:
|
|
# this is a gap
|
|
res.append(Segment(None, v))
|
|
|
|
if res[-1].id is None:
|
|
# random gap at the end? no thanks
|
|
res = res[:-1]
|
|
|
|
return res
|
|
|
|
|
|
def calc_checksum(files: list[Segment]) -> int:
|
|
acc = 0
|
|
pos = 0
|
|
for file in files:
|
|
if file.id is not None:
|
|
# sum of sequence of n consecutive integers is (n / 2)(first + last)
|
|
acc += int((file.len / 2) * ((pos * 2) + file.len - 1) * file.id)
|
|
pos += file.len
|
|
return acc
|
|
|
|
|
|
def one(instr: str):
|
|
files = parse(instr)
|
|
|
|
i = 0
|
|
while i < len(files) - 1:
|
|
f = files[i]
|
|
last_file = files[-1]
|
|
|
|
if last_file.id is None:
|
|
files = files[:-1]
|
|
continue
|
|
|
|
if f.id is None:
|
|
# _debug(i, files[i])
|
|
last_file = files[-1]
|
|
|
|
assert last_file.id is not None
|
|
if last_file.len > f.len:
|
|
files[-1] = Segment(last_file.id, last_file.len - f.len)
|
|
files[i] = Segment(last_file.id, f.len)
|
|
elif last_file.len == f.len:
|
|
# we're gonna move all of this so just delete it
|
|
files = files[:-1]
|
|
files[i] = last_file
|
|
else:
|
|
# TODO: when we haven't got enough in this last file so we need to split f into two
|
|
files = files[:-1]
|
|
files[i] = last_file
|
|
assert f.len - last_file.len != 0
|
|
files.insert(i+1, Segment(None, f.len - last_file.len))
|
|
|
|
i += 1
|
|
|
|
return calc_checksum(files)
|
|
|
|
|
|
def two(instr: str):
|
|
files = parse(instr)
|
|
|
|
gaps = defaultdict(list)
|
|
|
|
for i, file in enumerate(files):
|
|
if file.id is None and file.len != 0:
|
|
gaps[file.len].append(i)
|
|
|
|
moved = set()
|
|
|
|
i = len(files) - 1
|
|
while i > 0:
|
|
f = files[i]
|
|
|
|
if f.id is not None and f.id not in moved:
|
|
gap_loc = None
|
|
for j, v in enumerate(files):
|
|
if j >= i:
|
|
break
|
|
|
|
if v.id is None and v.len >= f.len:
|
|
gap_loc = j
|
|
break
|
|
|
|
if gap_loc is not None:
|
|
v = files[gap_loc]
|
|
if v.len == f.len:
|
|
files[i], files[gap_loc] = files[gap_loc], files[i]
|
|
else:
|
|
files[i] = Segment(None, f.len)
|
|
files[gap_loc] = f
|
|
files.insert(gap_loc + 1, Segment(None, v.len - f.len))
|
|
moved.add(f.id)
|
|
|
|
i -= 1
|
|
|
|
return calc_checksum(files)
|
|
|
|
|
|
def _debug(*args, **kwargs):
|
|
kwargs["file"] = sys.stderr
|
|
print(*args, **kwargs)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
if len(sys.argv) < 2 or sys.argv[1] not in ["1", "2"]:
|
|
print("Missing day argument", file=sys.stderr)
|
|
sys.exit(1)
|
|
inp = sys.stdin.read().strip()
|
|
if sys.argv[1] == "1":
|
|
print(one(inp))
|
|
else:
|
|
print(two(inp)) |