aoc/year2024/
day11.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
//! # Plutonian Pebbles
//!
//! The transformation rules have any interesting property that the total number of
//! distinct stone numbers is not very large, about 2000 for part one and 4000 for part two.
//!
//! This means that we can store the count of each distinct stone in a small contigous array
//! that is much faster to process than a recursive memoization approach.
use crate::util::hash::*;
use crate::util::parse::*;

pub fn parse(input: &str) -> Vec<u64> {
    input.iter_unsigned().collect()
}

pub fn part1(input: &[u64]) -> u64 {
    count(input, 25)
}

pub fn part2(input: &[u64]) -> u64 {
    count(input, 75)
}

fn count(input: &[u64], blinks: usize) -> u64 {
    // Allocate enough room to prevent reallocations.
    let mut stones = Vec::with_capacity(5000);
    // Maps number on stone to a much smaller contigous range of indices.
    let mut indices = FastMap::with_capacity(5000);
    // Numbers of any new stones generated during the previous blink.
    let mut todo = Vec::new();
    // Amount of each stone of a particular number.
    let mut current = Vec::new();

    // Initialize stones from input.
    for &number in input {
        if let Some(&index) = indices.get(&number) {
            current[index] += 1;
        } else {
            indices.insert(number, indices.len());
            todo.push(number);
            current.push(1);
        }
    }

    for _ in 0..blinks {
        // If a stone number has already been seen then return its index,
        // otherwise queue it for processing during the next blink.
        let numbers = todo;
        todo = Vec::with_capacity(200);

        let mut index_of = |number| {
            let size = indices.len();
            *indices.entry(number).or_insert_with(|| {
                todo.push(number);
                size
            })
        };

        // Apply the transformation logic to stones added in the previous blink.
        for number in numbers {
            let (first, second) = if number == 0 {
                (index_of(1), usize::MAX)
            } else {
                let digits = number.ilog10() + 1;
                if digits % 2 == 0 {
                    let power = 10_u64.pow(digits / 2);
                    (index_of(number / power), index_of(number % power))
                } else {
                    (index_of(number * 2024), usize::MAX)
                }
            };

            stones.push((first, second));
        }

        // Add amount of each stone to either 1 or 2 stones in the next blink.
        let mut next = vec![0; indices.len()];

        for (&(first, second), amount) in stones.iter().zip(current) {
            next[first] += amount;
            if second != usize::MAX {
                next[second] += amount;
            }
        }

        current = next;
    }

    current.iter().sum()
}