From 9c1fa8e84f34322b797bc0a4df4b79e036bd3b46 Mon Sep 17 00:00:00 2001
From: Daniel Lin <ephemient@gmail.com>
Date: Mon, 9 Dec 2024 03:51:44 -0500
Subject: [PATCH] Day 9: Disk Defragmenter

---
 README.md          |  2 +-
 py/aoc2024/day9.py | 65 ++++++++++++++++++++++++++++++++++++++++++++++
 py/pyproject.toml  |  1 +
 3 files changed, 67 insertions(+), 1 deletion(-)
 create mode 100644 py/aoc2024/day9.py

diff --git a/README.md b/README.md
index ffacf4e..353a185 100644
--- a/README.md
+++ b/README.md
@@ -13,4 +13,4 @@ Development occurs in language-specific directories:
 |[Day6.hs](hs/src/Day6.hs)|[Day6.kt](kt/aoc2024-lib/src/commonMain/kotlin/com/github/ephemient/aoc2024/Day6.kt)|[day6.py](py/aoc2024/day6.py)|[day6.rs](rs/src/day6.rs)|
 |[Day7.hs](hs/src/Day7.hs)|[Day7.kt](kt/aoc2024-lib/src/commonMain/kotlin/com/github/ephemient/aoc2024/Day7.kt)|[day7.py](py/aoc2024/day7.py)|[day7.rs](rs/src/day7.rs)|
 |[Day8.hs](hs/src/Day8.hs)|[Day8.kt](kt/aoc2024-lib/src/commonMain/kotlin/com/github/ephemient/aoc2024/Day8.kt)|[day8.py](py/aoc2024/day8.py)|[day8.rs](rs/src/day8.rs)|
-|[Day9.hs](hs/src/Day9.hs)|[Day9.kt](kt/aoc2024-lib/src/commonMain/kotlin/com/github/ephemient/aoc2024/Day9.kt)|||
+|[Day9.hs](hs/src/Day9.hs)|[Day9.kt](kt/aoc2024-lib/src/commonMain/kotlin/com/github/ephemient/aoc2024/Day9.kt)|[day9.py](py/aoc2024/day9.py)||
diff --git a/py/aoc2024/day9.py b/py/aoc2024/day9.py
new file mode 100644
index 0000000..94a49c5
--- /dev/null
+++ b/py/aoc2024/day9.py
@@ -0,0 +1,65 @@
+"""
+Day 9: Disk Defragmenter
+"""
+
+from itertools import accumulate
+
+SAMPLE_INPUT = """
+2333133121414131402
+"""
+
+
+def _rangesum(start: int, size: int) -> int:
+    return (2 * start + size - 1) * size // 2
+
+
+def part1(data: str) -> int:
+    """
+    >>> part1(SAMPLE_INPUT)
+    1928
+    """
+    chunks = [int(c) for c in data if c.isdigit()]
+    total, offset, i, j = 0, 0, 0, len(chunks) - 1
+    while i <= j:
+        if not i % 2:
+            size = chunks[i]
+            total += i // 2 * _rangesum(offset, size)
+            offset += size
+            i += 1
+        elif not j % 2:
+            size = min(chunks[i], chunks[j])
+            total += j // 2 * _rangesum(offset, size)
+            offset += size
+            chunks[i] -= size
+            if chunks[i] <= 0:
+                i += 1
+            chunks[j] -= size
+            if chunks[j] <= 0:
+                j -= 1
+        else:
+            j -= 1
+    return total
+
+
+def part2(data: str) -> int:
+    """
+    >>> part2(SAMPLE_INPUT)
+    2858
+    """
+    chunks = [int(c) for c in data if c.isdigit()]
+    offsets = list(accumulate(chunks, initial=0))
+    total = 0
+    for i in range(len(chunks) - 1 & ~1, -1, -2):
+        size = chunks[i]
+        offset = offsets[i]
+        for j in range(1, i, 2):
+            if chunks[j] >= size:
+                offset = offsets[j]
+                offsets[j] += size
+                chunks[j] -= size
+                break
+        total += i // 2 * _rangesum(offset, size)
+    return total
+
+
+parts = (part1, part2)
diff --git a/py/pyproject.toml b/py/pyproject.toml
index 896c306..a56485e 100644
--- a/py/pyproject.toml
+++ b/py/pyproject.toml
@@ -32,6 +32,7 @@ day5 = "aoc2024.day5:parts"
 day6 = "aoc2024.day6:parts"
 day7 = "aoc2024.day7:parts"
 day8 = "aoc2024.day8:parts"
+day9 = "aoc2024.day9:parts"
 
 [build-system]
 requires = ["poetry-core"]