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//! backwards 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.
23const TRIANGLE: [usize; 10] = [0, 0, 1, 3, 6, 10, 15, 21, 28, 36];
24
25/// Remove any trailing newlines and convert to `usize`.
26pub fn parse(input: &str) -> Vec<usize> {
27    input.trim().bytes().map(|b| (b - b'0') as usize).collect()
28}
29
30/// Block by block checksum comparison that doesn't allocate any memory.
31pub fn part1(disk: &[usize]) -> usize {
32    // Start at the first free block and the last file.
33    let mut left = 0;
34    let mut right = disk.len() - 2 + disk.len() % 2;
35    let mut needed = disk[right];
36    let mut block = 0;
37    let mut checksum = 0;
38
39    while left < right {
40        // When moving to the next free block, add the checksum for the file we're skipping over.
41        (checksum, block) = update(checksum, block, left, disk[left]);
42        let mut available = disk[left + 1];
43        left += 2;
44
45        while available > 0 {
46            if needed == 0 {
47                if left == right {
48                    break;
49                }
50                right -= 2;
51                needed = disk[right];
52            }
53
54            // Take as much space as possible from the current free block range.
55            let size = needed.min(available);
56            (checksum, block) = update(checksum, block, right, size);
57            available -= size;
58            needed -= size;
59        }
60    }
61
62    // Account for any remaining file blocks left over.
63    (checksum, _) = update(checksum, block, right, needed);
64    checksum
65}
66
67pub fn part2(disk: &[usize]) -> usize {
68    let mut block = 0;
69    let mut checksum = 0;
70    let mut free: Vec<_> = repeat_with(|| Vec::with_capacity(1_100)).take(10).collect();
71
72    // Build a min-heap (leftmost free block first) where the size of each block is
73    // implicit in the index of the array.
74    for (index, &size) in disk.iter().enumerate() {
75        if index % 2 == 1 && size > 0 {
76            free[size].push(block);
77        }
78
79        block += size;
80    }
81
82    // Add sentinel value and reverse vecs so that smallest blocks are last.
83    for heap in &mut free {
84        heap.push(block);
85        heap.reverse();
86    }
87
88    for (index, &size) in disk.iter().enumerate().rev() {
89        block -= size;
90
91        // Count any previous free blocks to decrement block offset correctly.
92        if index % 2 == 1 {
93            continue;
94        }
95
96        // Find the leftmost free block that can fit the file (if any).
97        let mut next_block = block;
98        let mut next_index = usize::MAX;
99
100        for (i, heap) in free.iter().enumerate().skip(size) {
101            let top = heap.len() - 1;
102            let first = heap[top];
103
104            if first < next_block {
105                next_block = first;
106                next_index = i;
107            }
108        }
109
110        // We can make smaller free block from bigger blocks but not the other way around.
111        // As an optimization if all blocks of the biggest size are after our position then
112        // we can ignore them.
113        if !free.is_empty() {
114            let biggest = free.len() - 1;
115            let top = free[biggest].len() - 1;
116
117            if free[biggest][top] > block {
118                free.pop();
119            }
120        }
121
122        // Update the checksum with the file's location (possibly unchanged).
123        let id = index / 2;
124        let extra = next_block * size + TRIANGLE[size];
125        checksum += id * extra;
126
127        // If we used a free block, remove then add back any leftover space.
128        if next_index != usize::MAX {
129            free[next_index].pop();
130
131            // Insert the new smaller block into the correct location.
132            // Most frequently this is directly at the end of the vector so even though this
133            // is technically `O(n)`, in practice it's faster than a real heap.
134            let to = next_index - size;
135            if to > 0 {
136                let mut i = free[to].len();
137                let value = next_block + size;
138
139                while free[to][i - 1] < value {
140                    i -= 1;
141                }
142
143                free[to].insert(i, value);
144            }
145        }
146    }
147
148    checksum
149}
150
151/// Convenience function to update checksum based on file location and size.
152#[inline]
153fn update(checksum: usize, block: usize, index: usize, size: usize) -> (usize, usize) {
154    let id = index / 2;
155    let extra = block * size + TRIANGLE[size];
156    (checksum + id * extra, block + size)
157}