Skip to main content

aoc/year2021/
day09.rs

1//! # Smoke Basin
2//!
3//! Part two is the classic [flood fill](https://en.wikipedia.org/wiki/Flood_fill) algorithm with a
4//! twist to return the size of the filled area. This algorithm can be implemented either as a
5//! [DFS](https://en.wikipedia.org/wiki/Depth-first_search) using recursion or as a
6//! [BFS](https://en.wikipedia.org/wiki/Breadth-first_search) using an auxiliary data structure
7//! such as a [`VecDeque`].
8//!
9//! This solution uses a DFS approach as it's faster and Rust's stack size limit seems enough
10//! to accommodate the maximum basin size. While we could use the [`Grid`] and [`Point`] modules
11//! to take in the original grid one line at a time with 2D coordinates, it turns out to be
12//! somewhat faster to instead just operate on a 1D array with an explicit border, where newline
13//! is treated the same as `'9'`, in order to eliminate bounds checking. We can also tweak the
14//! flood fill to track the lowest value seen along the way, to share the work between part
15//! one and part two.
16//!
17//! [`VecDeque`]: std::collections::VecDeque
18//! [`Grid`]: crate::util::grid
19//! [`Point`]: crate::util::point
20
21pub struct Basin {
22    lowest: u32, // Lowest integer seen within basin so far.
23    size: u32,   // Number of cells in the basin.
24}
25
26pub fn parse(input: &str) -> Vec<Basin> {
27    // Create a larger grid with all borders already filled with 9.
28    let width = input.lines().next().unwrap().len() + 1;
29    let mut grid = Vec::with_capacity(input.len() + 2 * width);
30    grid.resize(width, b'9');
31    grid.extend_from_slice(input.as_bytes());
32    grid.resize(grid.len() + width, b'9');
33
34    // Collect all basins in the grid. Masking with 15 turns '0' through '9' into their numeric
35    // value, and '\n' into 10, so that we can use newline as a second barrier character.
36    let mut basins = Vec::with_capacity(256);
37    for idx in width..width + input.len() {
38        if grid[idx] & 15 < 9 {
39            basins.push(flood_fill(&mut grid, idx, width as isize));
40        }
41    }
42
43    // Note that select_nth_unstable will partition the array faster than a full sort. With the
44    // partition in place, the final three elements are the largest.
45    let pivot = basins.len() - 3;
46    basins.select_nth_unstable_by_key(pivot, |b| b.size);
47
48    basins
49}
50
51pub fn part1(basins: &[Basin]) -> u32 {
52    basins.iter().map(|b| b.lowest + 1).sum()
53}
54
55pub fn part2(basins: &[Basin]) -> u32 {
56    // The list of basins is not sorted overall, but does have the largest three at the end.
57    basins[basins.len() - 3..].iter().map(|b| b.size).product()
58}
59
60fn flood_fill(grid: &mut [u8], idx: usize, width: isize) -> Basin {
61    let mut lowest = (grid[idx] & 15) as u32;
62    let mut size = 1;
63    grid[idx] = b'9';
64
65    for delta in [1, -1, width, -width] {
66        let other = idx.wrapping_add(delta as usize);
67        if grid[other] & 15 < 9 {
68            let basin = flood_fill(grid, other, width);
69            lowest = lowest.min(basin.lowest);
70            size += basin.size;
71        }
72    }
73
74    Basin { lowest, size }
75}