day07: part 2

diff --git a/src/day07.rs b/src/day07.rs
index 6d10e2466e35f0dd73abcd3416b062056c10adc6..d18fdab5ee88e9e86520e0434b067204dcc81cd1 100644 (file)
--- a/src/day07.rs
+++ b/src/day07.rs
@@ -28,7 +28,7 @@ impl Puzzle {
         beam.push((self.start.0 + 1, self.start.1));
         while let Some(pos) = beam.pop() {
             if self.map.contains_key(&pos) {
-                let mut location = self.map.get_mut(&pos).unwrap();
+                let location = self.map.get_mut(&pos).unwrap();
                 if location.1 {
                     continue;
                 }
@@ -46,11 +46,66 @@ impl Puzzle {
         }
         splitters
     }
+
+    fn part2(&mut self) -> u64 {
+        self.part1();
+        let first_pos = (self.start.0 + 1, self.start.1);
+        // keep track of all beams locations
+        let mut beams_pos: Vec<(usize, usize)> = self.map.iter()
+            .filter(|(_, v)| v.1 && v.0 == '.')
+            .map(|(k, _)| *k)
+            .collect();
+        beams_pos.sort_by(|a, b| a.cmp(&b));
+
+        // keep track of the number of timelines for each visited location
+        let mut beams_timelines: HashMap<(usize, usize), u64> = self.map.iter()
+            .filter(|(_, v)| v.1 && v.0 == '.')
+            .map(|(k, _)| {
+                if *k == first_pos {
+                    (*k, 1)
+                } else {
+                    (*k, 0)
+                }
+            })
+            .collect();
+
+        // compute the number of timelines per line
+        beams_pos.iter()
+            .for_each(|pos| {
+                let timeline = *beams_timelines.get(&pos).unwrap();
+                let next_pos = (pos.0 + 1, pos.1);
+                if self.map.contains_key(&next_pos) {
+                    let next_char = self.map.get_mut(&next_pos).unwrap().0; // it exists
+                    match next_char {
+                        '.' => {
+                            let next_timeline = beams_timelines.get_mut(&next_pos).unwrap();
+                            *next_timeline += timeline;
+                        },
+                        '^' => {
+                            [next_pos.1 - 1, next_pos.1 + 1].into_iter()
+                                .for_each(|j| {
+                                    let p = (next_pos.0, j);
+                                    let next_timeline = beams_timelines.get_mut(&p).unwrap();
+                                    *next_timeline += timeline;
+                                });
+                        },
+                        _ => unreachable!(),
+                    }
+                }
+            });
+
+        let last_line = beams_pos.iter().last().unwrap().0;
+        beams_timelines.iter()
+            .filter(|(k, _)| k.0 == last_line)
+            .map(|(_, v)| v)
+            .sum()
+    }
 }
 
 pub fn run(input: &str) -> Result<(), Box<dyn Error>> {
     let mut p = Puzzle::new(input);
     println!("part1: {}", p.part1());
+    println!("part2: {}", p.part2());
     Ok(())
 }
 
@@ -75,10 +130,56 @@ mod tests {
 ...............
 .^.^.^.^.^...^.
 ...............
+";
+
+// .......S.......
+// .......|.......
+// ......|^|......
+// ......|.|......
+// .....|^|^|.....
+// .....|.|.|.....
+    static TEST_TIMELINES_1: &str = "\
+.......S.......
+...............
+.......^.......
+...............
+......^.^......
+...............
+";
+
+// .......S.......
+// .......|.......
+// ......|^|......
+// ......|.|......
+// .....|^|^|.....
+// .....|.|.|.....
+// ....|^|^|^|....
+// ....|.|.|.|....
+    static TEST_TIMELINES_2: &str = "\
+.......S.......
+...............
+.......^.......
+...............
+......^.^......
+...............
+.....^.^.^.....
+...............
 ";
 
     #[test]
     fn test_part1() {
         assert_eq!(21, Puzzle::new(TEST_INPUT).part1());
     }
+    #[test]
+    fn test_part2() {
+        assert_eq!(40, Puzzle::new(TEST_INPUT).part2());
+    }
+    #[test]
+    fn test_part2_timelines_1() {
+        assert_eq!(4, Puzzle::new(TEST_TIMELINES_1).part2());
+    }
+    #[test]
+    fn test_part2_timelines_2() {
+        assert_eq!(8, Puzzle::new(TEST_TIMELINES_2).part2());
+    }
 }