1 files changed, 71 insertions, 0 deletions
diff --git a/09/src/part-2.rs b/09/src/part-2.rs
new file mode 100644
index 0000000..84e724c
--- /dev/null
+++ b/09/src/part-2.rs
@@ -0,0 +1,71 @@
+mod tensor;
+
+use std::io::{BufRead};
+
+pub type Matrix = tensor::Tensor<usize, 2>;
+pub type Index = tensor::Index<2>;
+
+pub fn main() {
+    let stdin = std::io::stdin();
+    let handle = stdin.lock();
+
+    let x: Matrix = {
+        let mut tmp: Vec<usize> = Vec::new();
+        let mut m: usize = 0;
+        for l in handle.lines() {
+            let line = l.unwrap();
+            tmp.extend(line.chars().map(|c| c.to_digit(10).unwrap() as usize));
+            m += 1;
+        }
+        let n = tmp.len()/m;
+        Matrix::new_from([m, n], tmp)
+    };
+
+    let mut neighbors: Vec<Option<& usize>> = Vec::new();
+    let mut low_points: Vec<Index> = Vec::new();
+    for i in 0..x.shape()[0] {
+        for j in 0..x.shape()[1] {
+            neighbors.clear();
+            neighbors.push(x.el(& [i, j.wrapping_sub(1)]));
+            neighbors.push(x.el(& [i, j+1]));
+            neighbors.push(x.el(& [i.wrapping_sub(1), j]));
+            neighbors.push(x.el(& [i+1, j]));
+
+            let min_neighbor: usize = (& neighbors).into_iter().filter_map(|& y| y).map(|& y| y).min().unwrap();
+            if min_neighbor > x[[i, j]] { low_points.push([i, j]); }
+        }
+    }
+
+    println!("Found {} low points.", low_points.len());
+
+    let mut basin_sizes: Vec<usize> = Vec::new();
+    let mut to_check: Vec<Index> = Vec::new();
+    let mut neighbor_idxs: Vec<Index> = Vec::new();
+    for & lp in (& low_points).into_iter() {
+        basin_sizes.push(0);
+        let mut visited: tensor::Tensor<bool, 2> = tensor::Tensor::new(*x.shape(), false);
+        to_check.push(lp);
+        //println!("Basin bottom: {},{} ({})", lp[0], lp[1], x[lp]);
+        while to_check.len() > 0 {
+            let p = to_check.pop().unwrap();
+            if !visited[p] {
+                visited[p] = true;
+                *basin_sizes.last_mut().unwrap() += 1;
+                neighbor_idxs.clear();
+                neighbor_idxs.push([p[0], p[1].wrapping_sub(1)]);
+                neighbor_idxs.push([p[0], p[1]+1]);
+                neighbor_idxs.push([p[0].wrapping_sub(1), p[1]]);
+                neighbor_idxs.push([p[0]+1, p[1]]);
+
+                (& neighbor_idxs).into_iter().filter(|& idx| x.in_bounds(idx))
+                    .filter(|& idx| !visited[idx])
+                    .filter(|& idx| x[idx] < 9 && x[idx] >= x[p])
+                    .map(|& idx| { to_check.push(idx); }).count();
+            }
+        }
+    }
+
+    basin_sizes.sort(); 
+    let basin_product: usize = basin_sizes.into_iter().rev().take(3).product();
+    println!("Product of the three largest basin sizes: {}", basin_product);
+}