From 7c021ab87f734007dd5d988cea9d1ef62056b434 Mon Sep 17 00:00:00 2001
From: Nettika <git@nettika.cat>
Date: Sun, 8 Dec 2024 23:22:34 -0800
Subject: [PATCH] Implement day 9 solver

---
 puzzles/7.py |   3 +-
 puzzles/9.py | 138 +++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 139 insertions(+), 2 deletions(-)
 create mode 100644 puzzles/9.py

diff --git a/puzzles/7.py b/puzzles/7.py
index 6406e8f..9437c31 100644
--- a/puzzles/7.py
+++ b/puzzles/7.py
@@ -1,7 +1,6 @@
 from functools import partial
-from typing import Callable, Iterator, NamedTuple
 from math import log
-
+from typing import Callable, Iterator, NamedTuple
 
 test_input = """
 190: 10 19
diff --git a/puzzles/9.py b/puzzles/9.py
new file mode 100644
index 0000000..eb84203
--- /dev/null
+++ b/puzzles/9.py
@@ -0,0 +1,138 @@
+from itertools import islice, pairwise
+from typing import NamedTuple
+
+from more_itertools import chunked, sliced
+
+test_input = """
+2333133121414131402
+"""
+
+test_solution_p1 = 1928
+test_solution_p2 = 2858
+
+
+def solve_p1(puzzle_input: str) -> int:
+    blocks = _parse_blocks(puzzle_input)
+    compacted_blocks = _compact_blocks(blocks)
+    return _hash_blocks(compacted_blocks)
+
+
+def solve_p2(puzzle_input: str) -> int:
+    blocks = _parse_blocks(puzzle_input)
+    compacted_blocks = _compact_blocks_no_fragmentation(blocks)
+    return _hash_blocks(compacted_blocks)
+
+
+class Block(NamedTuple):
+    id: int
+    position: int
+    size: int
+
+
+def _parse_blocks(puzzle_input: str) -> list[Block]:
+    blocks: list[Block] = []
+    cursor_position = 0
+    current_id = 0
+
+    snapshot = map(int, puzzle_input.strip())
+    for pair in chunked(snapshot, 2):
+        if data_size := pair[0]:
+            blocks.append(Block(current_id, cursor_position, data_size))
+            cursor_position += data_size
+            current_id += 1
+        if len(pair) == 2:
+            cursor_position += pair[1]
+
+    return blocks
+
+
+def _compact_blocks(blocks: list[Block]) -> list[Block]:
+    assert blocks
+    blocks = blocks.copy()
+
+    compacted_blocks: list[Block] = []
+    current_position = 0
+    moving_block = blocks.pop()
+
+    while blocks:
+        next_block = blocks.pop(0)
+        free_space = next_block.position - current_position
+        while free_space:
+            assert moving_block.size > 0
+            allocation = min(free_space, moving_block.size)
+            compacted_blocks.append(
+                Block(
+                    moving_block.id,
+                    current_position,
+                    allocation,
+                )
+            )
+            current_position += allocation
+            free_space -= allocation
+            if moving_block.size > allocation:
+                moving_block = Block(
+                    moving_block.id,
+                    moving_block.position + allocation,
+                    moving_block.size - allocation,
+                )
+            elif blocks:
+                moving_block = blocks.pop()
+            else:
+                break
+
+        compacted_blocks.append(next_block)
+        current_position += next_block.size
+
+    compacted_blocks.append(
+        Block(
+            moving_block.id,
+            current_position,
+            moving_block.size,
+        )
+    )
+
+    return compacted_blocks
+
+
+def _compact_blocks_no_fragmentation(blocks: list[Block]) -> list[Block]:
+    assert blocks
+    blocks = blocks.copy()
+
+    compacted_blocks: list[Block] = []
+
+    while len(blocks) > 1:
+        moving_block = blocks[-1]
+        for (_, a), (i, b) in pairwise(enumerate(blocks)):
+            free_space = b.position - a.position - a.size
+            if free_space >= moving_block.size:
+                modified_block = Block(
+                    moving_block.id,
+                    a.position + a.size,
+                    moving_block.size,
+                )
+                blocks.pop()
+                blocks.insert(i, modified_block)
+                break
+        else:
+            blocks.pop()
+            compacted_blocks.append(moving_block)
+
+    compacted_blocks.append(blocks[0])
+
+    return list(sorted(compacted_blocks, key=lambda block: block.position))
+
+
+def _print_blocks(blocks: list[Block]) -> str:
+    snapshot = []
+    for a, b in pairwise(blocks):
+        snapshot.append(str(a.id) * a.size)
+        snapshot.append("." * (b.position - a.position - a.size))
+    snapshot.append(str(blocks[-1].id) * blocks[-1].size)
+
+
+def _hash_blocks(blocks: list[Block]) -> int:
+    return sum(
+        block.id * pos
+        for block in blocks
+        for pos in range(block.position, block.position + block.size)
+    )