aboutsummaryrefslogtreecommitdiff
path: root/2024/09/puzzle-2.py
diff options
context:
space:
mode:
Diffstat (limited to '2024/09/puzzle-2.py')
-rwxr-xr-x2024/09/puzzle-2.py59
1 files changed, 59 insertions, 0 deletions
diff --git a/2024/09/puzzle-2.py b/2024/09/puzzle-2.py
new file mode 100755
index 0000000..7ae2a96
--- /dev/null
+++ b/2024/09/puzzle-2.py
@@ -0,0 +1,59 @@
+#!/usr/bin/python3
+
+
+import bisect
+import collections
+import operator
+
+
+Block = collections.namedtuple("Block", ["pos", "size"])
+File = collections.namedtuple("File", ["id", "pos", "size"])
+
+
+def first_free_block(free: list[list[int]], file: File) -> Block | None:
+ size = None
+ pos = file.pos
+ for i, xs in enumerate(free[file.size:], file.size):
+ if len(xs) > 0 and xs[0] < pos:
+ size = i
+ pos = xs[0]
+ return size and Block(pos, size)
+
+
+def main() -> None:
+ with open("input") as f:
+ s = f.read()
+
+ pos = 0
+ files = []
+ free = [[] for _ in range(10)]
+ for i, ch in enumerate(s):
+ size = int(ch)
+ if (i & 1) == 0:
+ files.append(File(i // 2, pos, size))
+ elif size != 0:
+ free[size].append(pos)
+ pos += size
+
+ posgetr = operator.attrgetter("pos")
+
+ i = len(files) - 1
+ while i >= 0:
+ if blk := first_free_block(free, files[i]):
+ f = files.pop(i)
+ free[blk.size].pop(0)
+ bisect.insort(files, File(f.id, blk.pos, f.size), key=posgetr)
+ if blk.size != f.size:
+ bisect.insort(free[blk.size - f.size], blk.pos + f.size)
+ else:
+ i -= 1
+
+ n = 0
+ for f in files:
+ for i in range(f.size):
+ n += f.id * (f.pos + i)
+ print(n)
+
+
+if __name__ == "__main__":
+ main() \ No newline at end of file