Year 2024 Day 9

Author: maneatingape

README.md

@@ -80,6 +80,7 @@ Performance is reasonable even on older hardware, for example a 2011 MacBook Pro
 | 6 | [Guard Gallivant](https://adventofcode.com/2024/day/6) | [Source](src/year2024/day06.rs) | 386 |
 | 7 | [Bridge Repair](https://adventofcode.com/2024/day/7) | [Source](src/year2024/day07.rs) | 136 |
 | 8 | [Resonant Collinearity](https://adventofcode.com/2024/day/8) | [Source](src/year2024/day08.rs) | 8 |
+| 9 | [Disk Fragmenter](https://adventofcode.com/2024/day/9) | [Source](src/year2024/day09.rs) | 163 |
 
 ## 2023
 

benches/benchmark.rs

@@ -300,4 +300,5 @@ mod year2024 {
     benchmark!(year2024, day06);
     benchmark!(year2024, day07);
     benchmark!(year2024, day08);
+    benchmark!(year2024, day09);
 }

src/lib.rs

@@ -299,4 +299,5 @@ pub mod year2024 {
     pub mod day06;
     pub mod day07;
     pub mod day08;
+    pub mod day09;
 }

src/main.rs

@@ -369,5 +369,6 @@ fn year2024() -> Vec<Solution> {
         solution!(year2024, day06),
         solution!(year2024, day07),
         solution!(year2024, day08),
+        solution!(year2024, day09),
     ]
 }

src/util/heap.rs

@@ -59,4 +59,9 @@ impl<K: Ord, V> MinHeap<K, V> {
     pub fn pop(&mut self) -> Option<(K, V)> {
         self.heap.pop().map(|w| (w.key, w.value))
     }
+
+    #[inline]
+    pub fn peek(&self) -> Option<(&K, &V)> {
+        self.heap.peek().map(|w| (&w.key, &w.value))
+    }
 }

src/year2024/day09.rs

@@ -0,0 +1,139 @@
+//! # Disk Fragmenter
+//!
+//! ## Part One
+//!
+//! Computes the checksum by simultaneously scanning forward for free blocks and
+//! backwards for files. No memory is allocated which makes it very fast.
+//!
+//! ## Part Two
+//!
+//! We build 10 [min heaps](https://en.wikipedia.org/wiki/Heap_(data_structure)) in an array to
+//! store the free space offsets. The index of the array implicitly stores the size of the
+//! free block.
+//!
+//! When moving a file to a free block, the corresponding heap is popped and then any leftover
+//! space is pushed back to the heap at a smaller index. The heap at index zero is not used
+//! but makes the indexing easier.
+use crate::util::heap::*;
+
+/// [Triangular numbers](https://en.wikipedia.org/wiki/Triangular_number) offset by two.
+/// Files can be a max size of 9 so we only need the first 10 values, including zero to make
+/// indexing easier.
+const EXTRA: [usize; 10] = [0, 0, 1, 3, 6, 10, 15, 21, 28, 36];
+
+/// Remove any trailing newlines and convert to `usize`.
+pub fn parse(input: &str) -> Vec<usize> {
+    input.trim().bytes().map(|b| (b - b'0') as usize).collect()
+}
+
+/// Block by block checksum comparison that doesn't allocate any memory.
+pub fn part1(disk: &[usize]) -> usize {
+    // Start at the first free block and the last file.
+    let mut free = 0;
+    let mut file = disk.len() + disk.len() % 2;
+
+    let mut available = 0;
+    let mut needed = 0;
+
+    let mut block = 0;
+    let mut checksum = 0;
+
+    while free < file {
+        // Take as much space as possible from the current free block range.
+        let size = needed.min(available);
+        (checksum, block) = update(checksum, block, file, size);
+        available -= size;
+        needed -= size;
+
+        // One or both of "available" and "free" could be zero.
+        if needed == 0 {
+            file -= 2;
+            needed = disk[file];
+        }
+
+        // When moving to the next free block, add the checksum for the file we're skipping over.
+        if available == 0 {
+            let size = disk[free];
+            (checksum, block) = update(checksum, block, free, size);
+            available = disk[free + 1];
+            free += 2;
+        }
+    }
+
+    // Account for any remaining file blocks left over.
+    (checksum, _) = update(checksum, block, file, needed);
+    checksum
+}
+
+pub fn part2(disk: &[usize]) -> usize {
+    let mut block = 0;
+    let mut checksum = 0;
+    let mut free: Vec<_> = (0..10).map(|_| MinHeap::with_capacity(1_000)).collect();
+
+    // Build a min-heap (leftmost free block first) where the size of each block is
+    // implicit in the index of the array.
+    for (index, &size) in disk.iter().enumerate() {
+        if index % 2 == 1 && size > 0 {
+            free[size].push(block, ());
+        }
+
+        block += size;
+    }
+
+    for (index, &size) in disk.iter().enumerate().rev() {
+        block -= size;
+
+        // Count any previous free blocks to decrement block offset correctly.
+        if index % 2 == 1 {
+            continue;
+        }
+
+        // Find the leftmost free block that can fit the file (if any).
+        let mut next_block = block;
+        let mut next_index = usize::MAX;
+
+        for (i, heap) in free.iter().enumerate().skip(size) {
+            if let Some((&first, ())) = heap.peek() {
+                if first < next_block {
+                    next_block = first;
+                    next_index = i;
+                }
+            }
+        }
+
+        // We can make smaller free block from bigger blocks but not the other way around.
+        // As an optimization if all blocks of the biggest size are after our position then
+        // we can ignore them.
+        if !free.is_empty() {
+            let last = free.len() - 1;
+            if let Some((&first, ())) = free[last].peek() {
+                if first > block {
+                    free.pop();
+                }
+            }
+        }
+
+        // Update the checksum with the file's location (possibly unchanged).
+        let id = index / 2;
+        let extra = next_block * size + EXTRA[size];
+        checksum += id * extra;
+
+        // If we used a free block, remove then add back any leftover space.
+        if next_index != usize::MAX {
+            free[next_index].pop();
+            if size < next_index {
+                free[next_index - size].push(next_block + size, ());
+            }
+        }
+    }
+
+    checksum
+}
+
+/// Convenience function to update checksum based on file location and size.
+#[inline]
+fn update(checksum: usize, block: usize, index: usize, size: usize) -> (usize, usize) {
+    let id = index / 2;
+    let extra = block * size + EXTRA[size];
+    (checksum + id * extra, block + size)
+}

tests/test.rs

@@ -289,4 +289,5 @@ mod year2024 {
     mod day06_test;
     mod day07_test;
     mod day08_test;
+    mod day09_test;
 }

tests/year2024/day09_test.rs

@@ -0,0 +1,15 @@
+use aoc::year2024::day09::*;
+
+const EXAMPLE: &str = "2333133121414131402";
+
+#[test]
+fn part1_test() {
+    let input = parse(EXAMPLE);
+    assert_eq!(part1(&input), 1928);
+}
+
+#[test]
+fn part2_test() {
+    let input = parse(EXAMPLE);
+    assert_eq!(part2(&input), 2858);
+}