109 lines
2.7 KiB
Rust
109 lines
2.7 KiB
Rust
use std::collections::HashSet;
|
|
|
|
pub fn parse_input(input: &str) -> Vec<Vec<u32>> {
|
|
input
|
|
.lines()
|
|
.map(|s| s.chars().map(|c| c.to_digit(10).unwrap()).collect())
|
|
.collect()
|
|
}
|
|
|
|
fn find_low_points(grid: &Vec<Vec<u32>>) -> Vec<(usize, usize)> {
|
|
let mut low_points = Vec::new();
|
|
let height = grid.len();
|
|
for i in 0..height {
|
|
let width = grid[i].len();
|
|
for j in 0..width {
|
|
if (i == 0 || grid[i - 1][j] > grid[i][j])
|
|
&& (i == height - 1 || grid[i + 1][j] > grid[i][j])
|
|
&& (j == 0 || grid[i][j - 1] > grid[i][j])
|
|
&& (j == width - 1 || grid[i][j + 1] > grid[i][j])
|
|
{
|
|
low_points.push((i, j));
|
|
}
|
|
}
|
|
}
|
|
return low_points;
|
|
}
|
|
|
|
pub fn part1(grid: Vec<Vec<u32>>) -> u32 {
|
|
let mut sum = 0;
|
|
for low_point in find_low_points(&grid) {
|
|
let depth = grid[low_point.0][low_point.1];
|
|
sum += depth + 1;
|
|
}
|
|
return sum;
|
|
}
|
|
|
|
fn find_basin_size(grid: &Vec<Vec<u32>>, location: (usize, usize)) -> u32 {
|
|
let mut explored_spots = HashSet::new();
|
|
let mut leaf_nodes = vec![location];
|
|
let height = grid.len();
|
|
let width = grid[0].len();
|
|
|
|
while leaf_nodes.len() > 0 {
|
|
let leaf_node = leaf_nodes.pop().unwrap();
|
|
explored_spots.insert(leaf_node);
|
|
|
|
let (i, j) = leaf_node;
|
|
if i > 0 && !explored_spots.contains(&(i - 1, j)) && grid[i - 1][j] != 9 {
|
|
leaf_nodes.push((i - 1, j));
|
|
}
|
|
if i < height - 1 && !explored_spots.contains(&(i + 1, j)) && grid[i + 1][j] != 9 {
|
|
leaf_nodes.push((i + 1, j));
|
|
}
|
|
if j > 0 && !explored_spots.contains(&(i, j - 1)) && grid[i][j - 1] != 9 {
|
|
leaf_nodes.push((i, j - 1));
|
|
}
|
|
if j < width - 1 && !explored_spots.contains(&(i, j + 1)) && grid[i][j + 1] != 9 {
|
|
leaf_nodes.push((i, j + 1));
|
|
}
|
|
}
|
|
|
|
return explored_spots.len() as u32;
|
|
}
|
|
|
|
fn find_basin_sizes(grid: &Vec<Vec<u32>>) -> Vec<u32> {
|
|
let mut sizes = Vec::new();
|
|
for low_point in find_low_points(&grid) {
|
|
sizes.push(find_basin_size(grid, low_point))
|
|
}
|
|
return sizes;
|
|
}
|
|
|
|
pub fn part2(grid: Vec<Vec<u32>>) -> u32 {
|
|
let mut basin_sizes = find_basin_sizes(&grid);
|
|
basin_sizes.sort_by(|a, b| b.cmp(a));
|
|
return basin_sizes[0] * basin_sizes[1] * basin_sizes[2];
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn part1_example() {
|
|
let input = vec![
|
|
vec![2, 1, 9, 9, 9, 4, 3, 2, 1, 0],
|
|
vec![3, 9, 8, 7, 8, 9, 4, 9, 2, 1],
|
|
vec![9, 8, 5, 6, 7, 8, 9, 8, 9, 2],
|
|
vec![8, 7, 6, 7, 8, 9, 6, 7, 8, 9],
|
|
vec![9, 8, 9, 9, 9, 6, 5, 6, 7, 8],
|
|
];
|
|
let result = part1(input);
|
|
assert_eq!(result, 15);
|
|
}
|
|
|
|
#[test]
|
|
fn part2_example() {
|
|
let input = vec![
|
|
vec![2, 1, 9, 9, 9, 4, 3, 2, 1, 0],
|
|
vec![3, 9, 8, 7, 8, 9, 4, 9, 2, 1],
|
|
vec![9, 8, 5, 6, 7, 8, 9, 8, 9, 2],
|
|
vec![8, 7, 6, 7, 8, 9, 6, 7, 8, 9],
|
|
vec![9, 8, 9, 9, 9, 6, 5, 6, 7, 8],
|
|
];
|
|
let result = part2(input);
|
|
assert_eq!(result, 1134);
|
|
}
|
|
}
|