Add 2024 day 9 part 2 solution

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