From 5ee4483e276f3a3fa0386cb368beb24a8acc4ec8 Mon Sep 17 00:00:00 2001
From: Rokas Puzonas <rokas.puz@gmail.com>
Date: Fri, 10 Dec 2021 14:46:22 +0200
Subject: [PATCH] feat: solve day 9 part 2

---
 src/day9.rs | 58 +++++++++++++++++++++++++++++++++++++++++++++++++----
 src/main.rs |  1 +
 2 files changed, 55 insertions(+), 4 deletions(-)

diff --git a/src/day9.rs b/src/day9.rs
index b8ceafa..72caa05 100644
--- a/src/day9.rs
+++ b/src/day9.rs
@@ -1,3 +1,5 @@
+use std::collections::HashSet;
+
 
 pub fn parse_input(input: &str) -> Vec<Vec<u32>> {
     input.lines()
@@ -5,7 +7,7 @@ pub fn parse_input(input: &str) -> Vec<Vec<u32>> {
         .collect()
 }
 
-fn find_low_points(grid: Vec<Vec<u32>>) -> Vec<u32> {
+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 {
@@ -15,7 +17,7 @@ fn find_low_points(grid: Vec<Vec<u32>>) -> Vec<u32> {
             && (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(grid[i][j]);
+                low_points.push((i, j));
             }
         }
     }
@@ -24,12 +26,47 @@ fn find_low_points(grid: Vec<Vec<u32>>) -> Vec<u32> {
 
 pub fn part1(grid: Vec<Vec<u32>>) -> u32 {
     let mut sum = 0;
-    for low_point in find_low_points(grid) {
-        sum += low_point + 1;
+    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 {
@@ -47,4 +84,17 @@ mod tests {
         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);
+    }
 }
diff --git a/src/main.rs b/src/main.rs
index 0b84156..8f005f4 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -40,6 +40,7 @@ fn run(day: i32, part: i32, input_filename: &str) {
         "8.1" => println!("{}", day8::part1(&day8::parse_input(&contents))),
         "8.2" => println!("{}", day8::part2(&day8::parse_input(&contents))),
         "9.1" => println!("{}", day9::part1(day9::parse_input(&contents))),
+        "9.2" => println!("{}", day9::part2(day9::parse_input(&contents))),
         _ => println!("Day {} part {} not found", day, part)
     }
 }