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

//! # Dumbo Octopus
//!
//! This puzzle resembles the [`Day 9`] flood fill a little. Since there are only 100 octopuses
//! a fixed size array is used both to track current energy levels and a second array to track
//! if an octopus has flashed this turn. Each time an octopus flashes it bumps its neighbors
//! energy levels, which can propagate recursively through the entire grid.
//!
//! [`Day 9`]: crate::year2021::day09

/// Pad the 10x10 grid by 1 on either side so that we can avoid boundary checks.
type Input = [u8; 144];
/// Offsets for each of the eight neighbors. The last four offsets will wrap around when
/// added to `u8` to give a smaller index.
const NEIGHBORS: [u8; 8] = [1, 11, 12, 13, 243, 244, 245, 255];

pub fn parse(input: &str) -> Input {
    let bytes: Vec<_> = input.lines().map(str::as_bytes).collect();
    let mut grid = [0; 144];

    for y in 0..10 {
        for x in 0..10 {
            grid[12 * (y + 1) + (x + 1)] = bytes[y][x] - b'0';
        }
    }

    grid
}

pub fn part1(input: &Input) -> usize {
    let (total, _) = simulate(input, |_, steps| steps < 100);
    total
}

pub fn part2(input: &Input) -> usize {
    let (_, steps) = simulate(input, |flashes, _| flashes < 100);
    steps
}

fn simulate(input: &Input, predicate: impl Fn(usize, usize) -> bool) -> (usize, usize) {
    let mut grid = *input;
    let mut flashed = [true; 144];
    let mut todo = Vec::with_capacity(100);

    let mut flashes = 0;
    let mut steps = 0;
    let mut total = 0;

    while predicate(flashes, steps) {
        flashes = 0;

        // Bump each octopus' energy level by one. If it flashes then add to `todo` queue.
        for y in 0..10 {
            for x in 0..10 {
                let index = 12 * (y + 1) + (x + 1);

                if grid[index] < 9 {
                    grid[index] += 1;
                    flashed[index] = false;
                } else {
                    grid[index] = 0;
                    flashed[index] = true;
                    todo.push(index as u8);
                }
            }
        }

        // Process each flash, possibly adding more to the queue.
        while let Some(index) = todo.pop() {
            flashes += 1;

            for offset in NEIGHBORS {
                let next = index.wrapping_add(offset) as usize;
                if flashed[next] {
                    continue;
                }

                if grid[next] < 9 {
                    grid[next] += 1;
                } else {
                    grid[next] = 0;
                    flashed[next] = true;
                    todo.push(next as u8);
                }
            }
        }

        steps += 1;
        total += flashes;
    }

    (total, steps)
}