Implement part 1 more efficiently via skips

Author: fwcd

day20/src/day20.rs

@@ -21,14 +21,14 @@ struct Racetrack {
 struct Cheat {
     start: Vec2<i32>,
     end: Option<Vec2<i32>>,
-    picos_left: i32,
+    picos_left: usize,
 }
 
 #[derive(Clone, Copy, Debug, Hash, PartialEq, Eq)]
 struct Node {
     pos: Vec2<i32>,
-    picos: i32,
-    cost: i32,
+    picos: usize,
+    cost: usize,
     cheat: Option<Cheat>,
 }
 
@@ -46,10 +46,21 @@ impl PartialOrd for Node {
 
 #[derive(Clone, Copy, Debug, Hash, PartialEq, Eq)]
 enum CheatPolicy {
-    Allowed { picos: i32 },
+    Allowed { picos: usize },
     Forbidden,
 }
 
+#[derive(Clone, Debug, Hash, PartialEq, Eq)]
+struct Skip {
+    positions: Vec<Vec2<i32>>,
+}
+
+impl Skip {
+    fn len(&self) -> usize {
+        self.positions.len()
+    }
+}
+
 impl Racetrack {
     fn height(&self) -> i32 {
         self.rows.len() as i32
@@ -67,8 +78,8 @@ impl Racetrack {
         (-1..=1)
             .flat_map(move |dy| (-1..=1)
                 .filter(move |&dx| (dx != 0) ^ (dy != 0))
-                .map(move |dx| Vec2::new(pos.x + dx, pos.y + dy)))
-            .filter(move |&neigh| self.in_bounds(neigh))
+                .map(move |dx| Vec2::new(pos.x + dx, pos.y + dy))
+                .filter(move |&neigh| self.in_bounds(neigh)))
     }
 
     fn locate(&self, c: char) -> Option<Vec2<i32>> {
@@ -78,7 +89,7 @@ impl Racetrack {
     }
 
     /// Traces out the racetrack, finding the distances to the given end.
-    fn find_distances_to_end(&self, start: Vec2<i32>, end: Vec2<i32>) -> HashMap<Vec2<i32>, i32> {
+    fn find_skips_to_end(&self, start: Vec2<i32>, end: Vec2<i32>) -> HashMap<Vec2<i32>, Skip> {
         let mut stack = Vec::new();
         let mut distances = HashMap::new();
 
@@ -92,60 +103,71 @@ impl Racetrack {
         }
 
         for (i, &pos) in stack.iter().enumerate() {
-            distances.insert(pos, (stack.len() - 1 - i) as i32);
+            distances.insert(pos, Skip { positions: stack[(i + 1)..].iter().cloned().collect() });
         }
 
         distances
     }
 
     /// Finds paths through the racetrack, ordered ascendingly by total picoseconds.
-    fn find_paths(&self, start: Vec2<i32>, end: Vec2<i32>, cheat_policy: CheatPolicy, condition: impl Fn(Node) -> bool) -> Vec<Node> {
+    fn count_paths(&self, start: Vec2<i32>, end: Vec2<i32>, cheat_policy: CheatPolicy, condition: impl Fn(usize) -> bool) -> i32 {
         // Your run-of-the-mill A* (Dijkstra + heuristic) implementation
 
+        let skips = self.find_skips_to_end(start, end);
         let mut queue = BinaryHeap::new();
         let mut visited = HashSet::new();
-        let mut paths = Vec::new();
+        let mut paths = 0;
 
         queue.push(Node { pos: start, picos: 0, cost: 0, cheat: None });
         visited.insert((start, None));
 
         while let Some(node) = queue.pop() {
+            let cheats_allowed = matches!(cheat_policy, CheatPolicy::Allowed { .. });
+            let can_cheat = cheats_allowed && node.cheat.map_or(true, |c| c.picos_left > 0);
+
             if node.pos == end {
-                if !condition(node) {
+                if !condition(node.picos) {
                     break;
                 }
-                println!("{node:?}");
-                paths.push(node);
+                // println!("{}", skips[&start].len() - node.cost);
+                paths += 1;
+                continue;
             }
 
-            for neigh in self.neighbors(node.pos) {
-                let is_wall = self[neigh] == '#';
-
-                let cheats_allowed = matches!(cheat_policy, CheatPolicy::Allowed { .. });
-                let can_cheat = cheats_allowed && node.cheat.map_or(true, |c| c.picos_left > 0);
-                let new_cheat = if let Some(cheat) = node.cheat {
-                    let is_ending = cheat.picos_left == 1;
-                    if is_ending && is_wall {
-                        continue;
-                    }
-                    Some(Cheat { start: cheat.start, end: if is_ending { Some(neigh) } else { cheat.end }, picos_left: (cheat.picos_left - 1).max(0) })
-                } else if cheats_allowed && is_wall {
-                    if let CheatPolicy::Allowed { picos: cheat_picos } = cheat_policy {
-                        Some(Cheat { start: node.pos, end: None, picos_left: cheat_picos })
+            if !can_cheat && skips.contains_key(&node.pos) {
+                // Skip the intermediate positions and hop straight to the end
+                let skip = &skips[&node.pos];
+                let new_picos = node.picos + skip.len();
+                queue.push(Node { pos: end, cost: new_picos, picos: new_picos, cheat: node.cheat });
+            } else {
+                // Search as usual
+                for neigh in self.neighbors(node.pos) {
+                    let is_wall = self[neigh] == '#';
+
+                    let new_cheat = if let Some(cheat) = node.cheat {
+                        let is_ending = cheat.picos_left == 1;
+                        if is_ending && is_wall {
+                            continue;
+                        }
+                        Some(Cheat { start: cheat.start, end: if is_ending { Some(neigh) } else { cheat.end }, picos_left: cheat.picos_left.max(1) - 1 })
+                    } else if cheats_allowed && is_wall {
+                        if let CheatPolicy::Allowed { picos: cheat_picos } = cheat_policy {
+                            Some(Cheat { start: node.pos, end: None, picos_left: cheat_picos - 1 })
+                        } else {
+                            unreachable!()
+                        }
                     } else {
-                        unreachable!()
-                    }
-                } else {
-                    node.cheat
-                };
+                        node.cheat
+                    };
 
-                if !visited.contains(&(neigh, new_cheat)) && (!is_wall || can_cheat) {
-                    visited.insert((neigh, new_cheat));
+                    if !visited.contains(&(neigh, new_cheat)) && (!is_wall || can_cheat) {
+                        visited.insert((neigh, new_cheat));
 
-                    let new_picos = node.picos + 1;
-                    let new_dist_to_end = (neigh.x.abs_diff(end.x) + neigh.y.abs_diff(end.y)) as i32;
-                    let new_cost = new_picos + new_dist_to_end;
-                    queue.push(Node { pos: neigh, picos: new_picos, cost: new_cost, cheat: new_cheat });
+                        let new_picos = node.picos + 1;
+                        let new_dist_to_end = (neigh.x.abs_diff(end.x) + neigh.y.abs_diff(end.y)) as usize;
+                        let new_cost = new_picos + new_dist_to_end;
+                        queue.push(Node { pos: neigh, picos: new_picos, cost: new_cost, cheat: new_cheat });
+                    }
                 }
             }
         }
@@ -175,14 +197,12 @@ fn main() {
     let start = track.locate('S').unwrap();
     let end = track.locate('E').unwrap();
 
-    let distances_to_end = track.find_distances_to_end(start, end);
-    let base_picos = distances_to_end[&start];
+    let skips = track.find_skips_to_end(start, end);
+    let base_picos = skips[&start].len();
 
-    let cheat_paths1 = track.find_paths(start, end, CheatPolicy::Allowed { picos: 2 }, |n| n.picos <= base_picos - 100);
-    let part1 = cheat_paths1.len();
+    let part1 = track.count_paths(start, end, CheatPolicy::Allowed { picos: 2 }, |picos| picos + 100 <= base_picos);
     println!("Part 1: {part1}");
 
-    let cheat_paths2 = track.find_paths(start, end, CheatPolicy::Allowed { picos: 20 }, |n| n.picos <= base_picos - 100);
-    let part2 = cheat_paths2.len();
+    let part2 = track.count_paths(start, end, CheatPolicy::Allowed { picos: 20 }, |picos| picos + 100 <= base_picos);
     println!("Part 2: {part2}");
 }