Clean up graph searching code.

day10/src/main.rs

@@ -1,45 +1,60 @@
-use std::collections::HashSet;
+use std::collections::{HashSet, VecDeque};
 
 fn main() {
     let input = std::fs::read_to_string("input").unwrap();
     let grid = Grid::new(&input);
     println!("{}", pt1(&grid));
+    println!("{}", pt2(&grid));
 }
 
 fn pt1(grid: &Grid) -> usize {
-    dfs(grid)
+    score(grid)
 }
 
-fn dfs(grid: &Grid) -> usize {
+fn pt2(grid: &Grid) -> usize {
+    rank(grid)
+}
+
+fn score(grid: &Grid) -> usize {
     let mut total = 0;
 
     for head in grid.heads.iter() {
-        let mut processed = HashSet::new();
+        let mut htotal = 0;
         let mut q = Vec::new();
-
+        let mut nines = HashSet::new();
         q.push(*head);
         while let Some(current) = q.pop() {
-            if !processed.contains(&current) {
+            if grid.get(&current).unwrap() == 9 && nines.insert(current) {
-                q.push(current);
+                htotal += 1;
             }
-
-            let mut do_top = true;
             let nexts = grid.next_step(&current);
             for next in nexts.iter() {
-                if !processed.contains(next) {
+                q.push(*next);
-                    q.push(*next);
-                    do_top = false;
-                }
-            }
-
-            if do_top {
-                let _ = q.pop();
-                processed.insert(current);
-                if grid.get(&current).unwrap() == 9 {
-                    total += 1;
-                }
             }
         }
+        total += htotal;
+    }
+
+    total
+}
+
+fn rank(grid: &Grid) -> usize {
+    let mut total = 0;
+
+    for head in grid.heads.iter() {
+        let mut htotal = 0;
+        let mut q = VecDeque::new();
+        q.push_back(*head);
+        while let Some(current) = q.pop_front() {
+            if grid.get(&current).unwrap() == 9 {
+                htotal += 1;
+            }
+            let nexts = grid.next_step(&current);
+            for next in nexts.iter() {
+                q.push_back(*next);
+            }
+        }
+        total += htotal;
     }
 
     total
@@ -128,6 +143,8 @@ impl Grid {
         out
     }
 
+    // returns list of neighbor nodes in a dag where the edges are all one step in
+    // distance
     fn next_step(&self, loc: &Loc) -> Vec<Loc> {
         let alt = self.rows[loc.row][loc.col];
         let target = alt + 1;
@@ -169,4 +186,10 @@ mod test {
         let g = Grid::new(INPUT);
         assert_eq!(36, pt1(&g));
     }
+
+    #[test]
+    fn p2() {
+        let g = Grid::new(INPUT);
+        assert_eq!(81, pt2(&g));
+    }
 }