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}