Add solution to day 9 part 2

day9/Cargo.lock

@@ -6,9 +6,25 @@ version = 3
 name = "day9"
 version = "0.1.0"
 dependencies = [
+ "itertools",
  "thiserror",
 ]
 
+[[package]]
+name = "either"
+version = "1.6.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e78d4f1cc4ae33bbfc157ed5d5a5ef3bc29227303d595861deb238fcec4e9457"
+
+[[package]]
+name = "itertools"
+version = "0.10.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a9a9d19fa1e79b6215ff29b9d6880b706147f16e9b1dbb1e4e5947b5b02bc5e3"
+dependencies = [
+ "either",
+]
+
 [[package]]
 name = "proc-macro2"
 version = "1.0.33"

day9/Cargo.toml

@@ -7,3 +7,4 @@ edition = "2021"
 
 [dependencies]
 thiserror = "1.0"
+itertools = "0.10"

day9/src/main.rs

@@ -1,4 +1,6 @@
 #![warn(clippy::all, clippy::pedantic)]
+use itertools::Itertools;
+use std::collections::{HashSet, VecDeque};
 use std::env;
 use std::fs::File;
 use std::io::{BufRead, BufReader};
@@ -10,55 +12,78 @@ enum Error {
     RowOutOfRange(usize),
 }
 
-fn find_low_point_indices(
+/// Get the indices (row, col) of all adjacent items that are in the input
+fn get_adjacent_indices(
     input: &[Vec<u32>],
-    depth: usize,
-) -> Result<impl Iterator<Item = usize> + '_, Error> {
+    (depth, col): (usize, usize),
+) -> Result<Vec<(usize, usize)>, Error> {
     let maybe_row = input.get(depth);
     if maybe_row.is_none() {
         return Err(Error::RowOutOfRange(depth));
     }
 
     let row = maybe_row.unwrap();
-    let iter = row.iter().enumerate().filter_map(move |(i, &item)| {
-        let adjacent = {
-            let mut res = vec![
-                row.get(i + 1),
-                input.get(depth + 1).map(|below_row| {
+    let mut res = vec![];
+
+    if col > 0 {
+        res.push((depth, col - 1));
+    }
+
+    if row.get(col + 1).is_some() {
+        res.push((depth, col + 1));
+    }
+
+    let below = input.get(depth + 1).map(|below_row| {
         below_row
-                        .get(i)
+            .get(col)
             .expect("below row does not match row length")
-                }),
-            ];
-
-            if i > 0 {
-                let left = row.get(i - 1);
-                res.push(left);
+    });
+    if below.is_some() {
+        res.push((depth + 1, col));
     }
-            if depth > 0 {
-                let above = input.get(depth - 1).map(|above_row| {
+
+    let have_above_item = || {
+        input
+            .get(depth - 1)
+            .map(|above_row| {
                 above_row
-                        .get(i)
+                    .get(col)
                     .expect("above row does not match row length")
-                });
-                res.push(above);
-            }
-
-            res
+            })
+            .is_some()
     };
-
-        if adjacent
-            .into_iter()
-            .filter(Option::is_some)
-            .all(|adjacent_item| item < *adjacent_item.unwrap())
-        {
-            Some(i)
-        } else {
-            None
+    if depth > 0 && have_above_item() {
+        res.push((depth - 1, col));
     }
-    });
 
-    Ok(iter)
+    Ok(res)
+}
+
+/// Get the indices of all low points in the input
+fn find_low_point_indices(input: &[Vec<u32>], depth: usize) -> Result<Vec<usize>, Error> {
+    let maybe_row = input.get(depth);
+    if maybe_row.is_none() {
+        return Err(Error::RowOutOfRange(depth));
+    }
+
+    let row = maybe_row.unwrap();
+    let mut res = vec![];
+    for (i, &item) in row.iter().enumerate() {
+        let adjacent_res = get_adjacent_indices(input, (depth, i));
+        if let Err(err) = adjacent_res {
+            return Err(err);
+        }
+
+        if adjacent_res
+            .unwrap()
+            .into_iter()
+            .all(|(adjacent_depth, adjacent_col)| item < input[adjacent_depth][adjacent_col])
+        {
+            res.push(i);
+        }
+    }
+
+    Ok(res)
 }
 
 fn part1(input: &[Vec<u32>]) -> u32 {
@@ -68,12 +93,54 @@ fn part1(input: &[Vec<u32>]) -> u32 {
         .map(|(depth, row)| {
             find_low_point_indices(input, depth)
                 .expect("failed to get row for depth analysis")
+                .into_iter()
                 .map(|idx| row[idx] + 1)
                 .sum::<u32>()
         })
         .sum()
 }
 
+fn part2(input: &[Vec<u32>]) -> u32 {
+    let basin_sizes = input.iter().enumerate().flat_map(|(depth, _)| {
+        let low_points =
+            find_low_point_indices(input, depth).expect("failed to get row for depth analysis");
+
+        low_points.into_iter().map(move |low_point_idx| {
+            let mut to_visit = [(depth, low_point_idx)]
+                .into_iter()
+                .collect::<VecDeque<_>>();
+            let mut visited = HashSet::<(usize, usize)>::new();
+            // 1 includes the low point
+            let mut num_in_basin = 1;
+
+            // Flood the board, terminating our search once we hit a nine
+            while !to_visit.is_empty() {
+                let (visiting_row, visiting_col) = to_visit.pop_front().unwrap();
+                let visiting = input[visiting_row][visiting_col];
+                let adjacent_iter = get_adjacent_indices(input, (visiting_row, visiting_col))
+                    .expect("failed to get adjacent items for bfs");
+                for (adjacent_row, adjacent_col) in adjacent_iter {
+                    if visited.contains(&(adjacent_row, adjacent_col)) {
+                        continue;
+                    }
+
+                    let adjacent = input[adjacent_row][adjacent_col];
+                    // flows from high to low, nine can't be part of the basin
+                    if adjacent > visiting && adjacent != 9 {
+                        num_in_basin += 1;
+                        to_visit.push_back((adjacent_row, adjacent_col));
+                        visited.insert((adjacent_row, adjacent_col));
+                    }
+                }
+            }
+
+            num_in_basin
+        })
+    });
+
+    basin_sizes.sorted().rev().take(3).product()
+}
+
 fn main() {
     let input_file_name = env::args().nth(1).expect("No input filename specified");
     let input_file = File::open(input_file_name).expect("Could not open input file");
@@ -91,4 +158,5 @@ fn main() {
         .collect::<Vec<_>>();
 
     println!("Part 1: {}", part1(&input_lines));
+    println!("Part 2: {}", part2(&input_lines));
 }