Skip to main content

aoc/year2024/
day09.rs

1//! # Disk Fragmenter
2//!
3//! ## Part One
4//!
5//! Computes the checksum by simultaneously scanning forward for free blocks and
6//! backward for files. No memory is allocated which makes it very fast.
7//!
8//! ## Part Two
9//!
10//! We build 10 [min heaps](https://en.wikipedia.org/wiki/Heap_(data_structure)) in an array to
11//! store the free space offsets. The index of the array implicitly stores the size of the
12//! free block. The heaps are implemented as a simple reversed `vec`. Usually items are added
13//! directly to the top of the heap, so this is faster than a real heap.
14//!
15//! When moving a file to a free block, the corresponding heap is popped and then any leftover
16//! space is pushed back to the heap at a smaller index. The heap at index zero is not used
17//! but makes the indexing easier.
18use std::iter::repeat_with;
19
20/// [Triangular numbers](https://en.wikipedia.org/wiki/Triangular_number) offset by two.
21/// Files can be a max size of 9 so we only need the first 10 values, including zero to make
22/// indexing easier.
23use crate::util::parse::*;
24
25const TRIANGLE: [usize; 10] = [0, 0, 1, 3, 6, 10, 15, 21, 28, 36];
26
27/// Remove any trailing newlines and convert to `usize`.
28pub fn parse(input: &str) -> Vec<usize> {
29    input.trim().bytes().map(u8::to_decimal).collect()
30}
31
32/// Block by block checksum comparison that doesn't allocate any memory.
33pub fn part1(disk: &[usize]) -> usize {
34    // Start at the first free block and the last file.
35    let mut left = 0;
36    let mut right = disk.len() - 2 + disk.len() % 2;
37    let mut needed = disk[right];
38    let mut block = 0;
39    let mut checksum = 0;
40
41    while left < right {
42        // When moving to the next free block, add the checksum for the file we're skipping over.
43        (checksum, block) = update(checksum, block, left, disk[left]);
44        let mut available = disk[left + 1];
45        left += 2;
46
47        while available > 0 {
48            if needed == 0 {
49                if left == right {
50                    break;
51                }
52                right -= 2;
53                needed = disk[right];
54            }
55
56            // Take as much space as possible from the current free block range.
57            let size = needed.min(available);
58            (checksum, block) = update(checksum, block, right, size);
59            available -= size;
60            needed -= size;
61        }
62    }
63
64    // Account for any remaining file blocks left over.
65    (checksum, _) = update(checksum, block, right, needed);
66    checksum
67}
68
69pub fn part2(disk: &[usize]) -> usize {
70    let mut block = 0;
71    let mut checksum = 0;
72    let mut free: Vec<_> = repeat_with(|| Vec::with_capacity(1_100)).take(10).collect();
73
74    // Build a min-heap (leftmost free block first) where the size of each block is
75    // implicit in the index of the array.
76    for (index, &size) in disk.iter().enumerate() {
77        if !index.is_multiple_of(2) && size > 0 {
78            free[size].push(block);
79        }
80
81        block += size;
82    }
83
84    // Add sentinel value and reverse vecs so that smallest blocks are last.
85    for heap in &mut free {
86        heap.push(block);
87        heap.reverse();
88    }
89
90    for (index, &size) in disk.iter().enumerate().rev() {
91        block -= size;
92
93        // Count any previous free blocks to decrement block offset correctly.
94        if !index.is_multiple_of(2) {
95            continue;
96        }
97
98        // Find the leftmost free block that can fit the file (if any).
99        let mut next_block = block;
100        let mut next_index = usize::MAX;
101
102        for (i, heap) in free.iter().enumerate().skip(size) {
103            let first = *heap.last().unwrap();
104
105            if first < next_block {
106                next_block = first;
107                next_index = i;
108            }
109        }
110
111        // We can make smaller free blocks from bigger blocks but not the other way around.
112        // As an optimization if all blocks of the biggest size are after our position then
113        // we can ignore them.
114        if free.last().is_some_and(|h| *h.last().unwrap() > block) {
115            free.pop();
116        }
117
118        // Update the checksum with the file's location (possibly unchanged).
119        let id = index / 2;
120        let extra = next_block * size + TRIANGLE[size];
121        checksum += id * extra;
122
123        // If we used a free block, remove then add back any leftover space.
124        if next_index != usize::MAX {
125            free[next_index].pop();
126
127            // Insert the new smaller block into the correct location.
128            // Most frequently this is directly at the end of the vector so even though this
129            // is technically `O(n)`, in practice it's faster than a real heap.
130            let to = next_index - size;
131            if to > 0 {
132                let mut i = free[to].len();
133                let value = next_block + size;
134
135                while free[to][i - 1] < value {
136                    i -= 1;
137                }
138
139                free[to].insert(i, value);
140            }
141        }
142    }
143
144    checksum
145}
146
147/// Convenience function to update checksum based on file location and size.
148#[inline]
149fn update(checksum: usize, block: usize, index: usize, size: usize) -> (usize, usize) {
150    let id = index / 2;
151    let extra = block * size + TRIANGLE[size];
152    (checksum + id * extra, block + size)
153}