151 lines
3.5 KiB
Rust
151 lines
3.5 KiB
Rust
use std::collections::HashSet;
|
|
|
|
use color_eyre::eyre;
|
|
use eyre::eyre;
|
|
|
|
use crate::lib::{IterExt, Vec2};
|
|
|
|
|
|
#[derive(Debug)]
|
|
struct Cell {
|
|
height: u32,
|
|
min: bool,
|
|
}
|
|
|
|
struct NeighborsIter {
|
|
x: usize,
|
|
y: usize,
|
|
rows: usize,
|
|
cols: usize,
|
|
state: usize,
|
|
}
|
|
|
|
impl Iterator for NeighborsIter {
|
|
type Item = (usize, usize);
|
|
fn next(&mut self) -> Option<Self::Item> {
|
|
let neighbor = match self.state {
|
|
0 if self.x > 0 => Some((self.y, self.x - 1)),
|
|
1 if self.y > 0 => Some((self.y - 1, self.x)),
|
|
2 if self.x < self.cols - 1 => Some((self.y, self.x + 1)),
|
|
3 if self.y < self.rows - 1 => Some((self.y + 1, self.x)),
|
|
_ => None,
|
|
};
|
|
self.state += 1;
|
|
|
|
if self.state > 4 {
|
|
None
|
|
}
|
|
else {
|
|
neighbor.or_else(|| self.next())
|
|
}
|
|
}
|
|
}
|
|
|
|
|
|
// struct CellWindow {
|
|
// before: Option<Cell>,
|
|
// cell: Cell,
|
|
// above: Option<Cell>,
|
|
// }
|
|
|
|
|
|
fn load(data: &str) -> eyre::Result<Vec2<Cell>> {
|
|
let num_columns = data.lines().next().unwrap().len();
|
|
let mut heights = Vec2::new(num_columns);
|
|
for line in data.lines() {
|
|
let mut row = Vec::new();
|
|
for c in line.chars() {
|
|
match c.to_digit(10) {
|
|
Some(d) => row.push(Cell {height: d, min: false}),
|
|
None => return Err(eyre!("Invalid character (not a digit): {}", c)),
|
|
}
|
|
}
|
|
heights.push(row);
|
|
}
|
|
Ok(heights)
|
|
}
|
|
|
|
|
|
fn scan(map: &mut Vec2<Cell>) {
|
|
for row in 0..map.row_count() {
|
|
for col in 0..map.col_count() {
|
|
let mut lt_left = true;
|
|
if col > 0 {
|
|
if map[row][col].height < map[row][col - 1].height {
|
|
map[row][col - 1].min = false;
|
|
}
|
|
else {
|
|
lt_left = false;
|
|
}
|
|
}
|
|
|
|
let mut lt_top = true;
|
|
if row > 0 {
|
|
if map[row][col].height < map[row - 1][col].height {
|
|
map[row - 1][col].min = false;
|
|
}
|
|
else {
|
|
lt_top = false;
|
|
}
|
|
}
|
|
|
|
if lt_left && lt_top {
|
|
map[row][col].min = true;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
|
|
fn low_points(map: &Vec2<Cell>) -> impl Iterator<Item=(usize, &Cell)> {
|
|
map.values().enumerate().filter(|(_, c)| c.min)
|
|
}
|
|
|
|
|
|
fn neighbors(map: &Vec2<Cell>, coords: (usize, usize)) -> NeighborsIter {
|
|
NeighborsIter {
|
|
x: coords.1,
|
|
y: coords.0,
|
|
rows: map.row_count(),
|
|
cols: map.col_count(),
|
|
state: 0,
|
|
}
|
|
}
|
|
|
|
|
|
fn basin_size (map: &Vec2<Cell>, start: (usize, usize), visited: &mut HashSet<(usize, usize)>) -> u32 {
|
|
let mut total = 1;
|
|
for coords in neighbors(map, start) {
|
|
if map[coords.0][coords.1].height < 9 && !visited.contains(&coords) {
|
|
visited.insert(coords);
|
|
total += basin_size(map, coords, visited);
|
|
}
|
|
}
|
|
total
|
|
}
|
|
|
|
|
|
fn basins(map: &Vec2<Cell>) -> u32 {
|
|
low_points(map)
|
|
.map(|(i, _)| {
|
|
let start = map.coords(i);
|
|
basin_size(map, start, &mut HashSet::from([start]))
|
|
})
|
|
.max_n(3)
|
|
.unwrap()
|
|
.iter()
|
|
.product()
|
|
}
|
|
|
|
|
|
pub fn run(data: &str) -> eyre::Result<(u32, u32)> {
|
|
// let sample = "2199943210\n3987894921\n9856789892\n8767896789\n9899965678";
|
|
let mut map = load(data)?;
|
|
scan(&mut map);
|
|
|
|
let one = low_points(&map).fold(0, |acc, (_, c)| acc + c.height + 1);
|
|
let two = basins(&map);
|
|
|
|
Ok((one, two))
|
|
}
|