aoc/year2024/
day09.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
//! # 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. The heaps are implemented as a simple reversed `vec`. Usually items are added
//! directly to the top of the heap, so this is faster than a real heap.
//!
//! 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 std::iter::repeat_with;

/// [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 TRIANGLE: [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 left = 0;
    let mut right = disk.len() - 2 + disk.len() % 2;
    let mut needed = disk[right];
    let mut block = 0;
    let mut checksum = 0;

    while left < right {
        // When moving to the next free block, add the checksum for the file we're skipping over.
        (checksum, block) = update(checksum, block, left, disk[left]);
        let mut available = disk[left + 1];
        left += 2;

        while available > 0 {
            if needed == 0 {
                if left == right {
                    break;
                }
                right -= 2;
                needed = disk[right];
            }

            // Take as much space as possible from the current free block range.
            let size = needed.min(available);
            (checksum, block) = update(checksum, block, right, size);
            available -= size;
            needed -= size;
        }
    }

    // Account for any remaining file blocks left over.
    (checksum, _) = update(checksum, block, right, needed);
    checksum
}

pub fn part2(disk: &[usize]) -> usize {
    let mut block = 0;
    let mut checksum = 0;
    let mut free: Vec<_> = repeat_with(|| Vec::with_capacity(1_100)).take(10).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;
    }

    // Add sentinel value and reverse vecs so that smallest blocks are last.
    for heap in &mut free {
        heap.push(block);
        heap.reverse();
    }

    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) {
            let top = heap.len() - 1;
            let first = heap[top];

            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 biggest = free.len() - 1;
            let top = free[biggest].len() - 1;

            if free[biggest][top] > block {
                free.pop();
            }
        }

        // Update the checksum with the file's location (possibly unchanged).
        let id = index / 2;
        let extra = next_block * size + TRIANGLE[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();

            // Insert the new smaller block into the correct location.
            // Most frequently this is directly at the end of the vector so even though this
            // is technically `O(n)`, in practice it's faster than a real heap.
            let to = next_index - size;
            if to > 0 {
                let mut i = free[to].len();
                let value = next_block + size;

                while free[to][i - 1] < value {
                    i -= 1;
                }

                free[to].insert(i, value);
            }
        }
    }

    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 + TRIANGLE[size];
    (checksum + id * extra, block + size)
}