day10, part2

day10/src/main.rs

@@ -1,19 +1,25 @@
-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)
 }
 
+fn pt2(grid: &Grid) -> usize {
+    bfs(grid)
+}
+
 fn dfs(grid: &Grid) -> usize {
     let mut total = 0;
 
     for head in grid.heads.iter() {
+        let mut htotal = 0;
         let mut processed = HashSet::new();
         let mut q = Vec::new();
 
@@ -34,12 +40,35 @@ fn dfs(grid: &Grid) -> usize {
 
             if do_top {
                 let _ = q.pop();
-                processed.insert(current);
+                if processed.insert(current) && grid.get(&current).unwrap() == 9 {
-                if grid.get(&current).unwrap() == 9 {
+                    htotal += 1;
-                    total += 1;
                 }
             }
         }
+        total += htotal;
+    }
+
+    total
+}
+
+fn bfs(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() {
+            let nexts = grid.next_step(&current);
+            for next in nexts.iter() {
+                if grid.get(next).unwrap() == 9 {
+                    htotal += 1;
+                }
+
+                q.push_back(*next);
+            }
+        }
+        total += htotal;
     }
 
     total
@@ -169,4 +198,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));
+    }
 }