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}