AoC21/src/day09.rs

92 lines
2.6 KiB
Rust

use itertools::Itertools;
use std::collections::{HashSet, VecDeque};
pub fn day9() {
// let input = r#"
// 2199943210
// 3987894921
// 9856789892
// 8767896789
// 9899965678
// "#;
let input = include_str!("../input/day9.txt");
let grid = input
.trim()
.lines()
.map(|l| {
l.trim()
.chars()
.map(|c| c.to_digit(10).unwrap())
.collect::<Vec<_>>()
})
.collect::<Vec<_>>();
let grid_slice = &grid.iter().map(|row| row.as_slice()).collect::<Vec<_>>();
let minima = find_local_minima(grid_slice);
let sum: u32 = minima.iter().map(|(row, col)| &grid[*row][*col] + 1).sum();
println!("p1 score sum: {}", sum);
dbg!(find_basin_size(&(0, 1), grid_slice));
dbg!(find_basin_size(&(0, 9), grid_slice));
let basins_product: i32 = minima
.iter()
.map(|pos| find_basin_size(pos, grid_slice) as i32)
.sorted_by_key(|count| -*count)
.take(3)
.product();
println!("p2 basin score product: {}", basins_product);
}
fn find_local_minima(grid: &[&[u32]]) -> Vec<(usize, usize)> {
let mut res = vec![];
let rows = grid.len();
let cols = grid[0].len();
for row in 0..rows {
for col in 0..cols {
let cur = grid[row][col];
let up = row == 0 || grid[row - 1][col] > cur;
let down = row == rows - 1 || grid[row + 1][col] > cur;
let left = col == 0 || grid[row][col - 1] > cur;
let right = col == cols - 1 || grid[row][col + 1] > cur;
if up && down && left && right {
res.push((row, col));
}
}
}
res
}
fn find_basin_size(pos: &(usize, usize), grid: &[&[u32]]) -> usize {
let rows = grid.len();
let cols = grid[0].len();
let mut queue = VecDeque::from([*pos]);
let mut visited: HashSet<(usize, usize)> = HashSet::new();
while let Some(p) = queue.pop_front() {
let height = grid[p.0][p.1];
if visited.contains(&p) || height == 9 {
continue;
}
visited.insert(p);
// println!("({}, {})", p.0, p.1);
if p.0 > 0 && grid[p.0 - 1][p.1] >= height {
queue.insert(0, (p.0 - 1, p.1));
}
if p.0 < rows - 1 && grid[p.0 + 1][p.1] >= height {
queue.insert(0, (p.0 + 1, p.1));
}
if p.1 > 0 && grid[p.0][p.1 - 1] >= height {
queue.insert(0, (p.0, p.1 - 1));
}
if p.1 < cols - 1 && grid[p.0][p.1 + 1] >= height {
queue.insert(0, (p.0, p.1 + 1));
}
}
visited.len()
}