diff --git a/aoc2021/src/day09.rs b/aoc2021/src/day09.rs index 3dd7c4c..9d21357 100644 --- a/aoc2021/src/day09.rs +++ b/aoc2021/src/day09.rs @@ -1,3 +1,4 @@ +use std::collections::HashSet; use std::fmt::Write; use anyhow::{Context, Result}; @@ -8,6 +9,7 @@ pub fn run() -> Result { let mut res = String::with_capacity(128); writeln!(res, "part 1: {}", part1(INPUT)?)?; + writeln!(res, "part 2: {}", part2(INPUT)?)?; Ok(res) } @@ -21,6 +23,25 @@ fn part1(input: &str) -> Result { .sum()) } +fn part2(input: &str) -> Result { + let mut height_map: HeightMap = input.parse()?; + + let low_points: Vec<_> = height_map.low_points().collect(); + + let mut bassin_sizes: Vec<_> = low_points + .iter() + .map(|&(x, y)| height_map.fill_basin(x, y)) + .collect(); + bassin_sizes.sort_unstable(); + + bassin_sizes + .iter() + .copied() + .skip(bassin_sizes.len() - 3) + .reduce(|acc, elem| acc * elem) + .context("couldn't find 3 bassins") +} + #[derive(Clone, Copy)] enum Neighbour { Up, @@ -53,9 +74,28 @@ struct HeightMap { heights: Vec, width: usize, height: usize, + filled_points: HashSet<(usize, usize)>, } impl HeightMap { + fn fill_basin(&mut self, x: usize, y: usize) -> u64 { + if self.get(x, y) == 9 { + return 0; + } + + if self.filled_points.contains(&(x, y)) { + return 0; + } + self.filled_points.insert((x, y)); + + let neighbours: Vec<_> = self.neighbours(x, y).collect(); + neighbours + .iter() + .map(|&(nx, ny)| self.fill_basin(nx, ny)) + .sum::() + + 1 + } + fn low_points(&self) -> impl Iterator + '_ { (0..self.height) .flat_map(|y| (0..self.width).map(move |x| (x, y))) @@ -111,6 +151,7 @@ impl std::str::FromStr for HeightMap { heights, width: width.context("0 lines parsed, width never computed")?, height, + filled_points: HashSet::new(), }) } } @@ -130,4 +171,14 @@ mod tests { fn part1_real() { assert_eq!(part1(INPUT).unwrap(), 522); } + + #[test] + fn part2_provided() { + assert_eq!(part2(PROVIDED).unwrap(), 1134); + } + + #[test] + fn part2_real() { + assert_eq!(part2(INPUT).unwrap(), 916688); + } }