part2: works up to the 4th iteration, then returns too high values

I think the issue is not testing both possible paths for direction keypads, but doing this naively blows up the compute complexity (and I got the wrong answer)

src/day21.rs

@@ -1,13 +1,23 @@
 use aoc_runner_derive::aoc;
-use itertools::Itertools;
+use rustc_hash::FxHashMap;
 use std::iter::repeat_n;
 
-trait KeypadRobot {
-    fn new() -> Self;
-    fn press(&mut self, target: u8) -> Vec<Vec<u8>>;
+#[derive(Clone, Debug)]
+enum KeypadRobot {
+    Number(NumberKeypadRobot),
+    Direction(DirectionKeypadRobot),
 }
 
-#[derive(Clone, Copy, Debug)]
+impl KeypadRobot {
+    fn press(&mut self, target: u8) -> Vec<FxHashMap<Vec<u8>, usize>> {
+        match self {
+            Self::Number(r) => r.press(target),
+            Self::Direction(r) => r.press(target),
+        }
+    }
+}
+
+#[derive(Clone, Debug)]
 struct NumberKeypadRobot {
     pointing_at: u8,
 }
@@ -31,31 +41,39 @@ impl NumberKeypadRobot {
         }
     }
 }
-impl KeypadRobot for NumberKeypadRobot {
+impl NumberKeypadRobot {
     fn new() -> Self {
         Self { pointing_at: b'A' }
     }
-    fn press(&mut self, target: u8) -> Vec<Vec<u8>> {
+    fn press(&mut self, target: u8) -> Vec<FxHashMap<Vec<u8>, usize>> {
         let cur_pos = Self::pos_of(self.pointing_at);
         let goal_pos = Self::pos_of(target);
         let x_ofs = goal_pos.0 - cur_pos.0;
         let y_ofs = goal_pos.1 - cur_pos.1;
 
         let mut paths = Vec::new();
-
+        // NOTE: no need to consider zig-zags since those paths will always require more button presses going back and forth
         if (cur_pos.0 + x_ofs, cur_pos.1) != Self::pos_of(b'X') {
             let mut x_first = Vec::new();
             x_first.extend(repeat_n(if x_ofs > 0 { b'>' } else { b'<' }, x_ofs.abs() as usize));
             x_first.extend(repeat_n(if y_ofs > 0 { b'v' } else { b'^' }, y_ofs.abs() as usize));
             x_first.push(b'A');
-            paths.push(x_first);
+            paths.push({
+                let mut h = FxHashMap::default();
+                h.insert(x_first, 1);
+                h
+            });
         }
         if (cur_pos.0, cur_pos.1 + y_ofs) != Self::pos_of(b'X') {
             let mut y_first = Vec::new();
             y_first.extend(repeat_n(if y_ofs > 0 { b'v' } else { b'^' }, y_ofs.abs() as usize));
             y_first.extend(repeat_n(if x_ofs > 0 { b'>' } else { b'<' }, x_ofs.abs() as usize));
             y_first.push(b'A');
-            paths.push(y_first);
+            paths.push({
+                let mut h = FxHashMap::default();
+                h.insert(y_first, 1);
+                h
+            });
         }
         if paths.is_empty() {
             panic!("all paths lead to the void");
@@ -66,13 +84,14 @@ impl KeypadRobot for NumberKeypadRobot {
     }
 }
 
-#[derive(Clone, Copy, Debug)]
-struct DirectionKeypadRobot<T: KeypadRobot> {
+#[derive(Clone, Debug)]
+struct DirectionKeypadRobot {
     pointing_at: u8,
-    child: Option<T>,
+    child: Option<Box<KeypadRobot>>,
+    id: usize,
 }
 
-impl<T: KeypadRobot> DirectionKeypadRobot<T> {
+impl DirectionKeypadRobot {
     fn pos_of(target: u8) -> (i8, i8) {
         match target {
             b'X' => (0, 0),
@@ -117,23 +136,31 @@ impl<T: KeypadRobot> DirectionKeypadRobot<T> {
         self.pointing_at = prev_point;
         path
     }
-}
-impl<T: KeypadRobot> KeypadRobot for DirectionKeypadRobot<T> {
-    fn new() -> Self {
+
+    fn new(id: usize, child: Option<Box<KeypadRobot>>) -> Self {
         Self {
+            id,
             pointing_at: b'A',
-            child: None,
+            child,
         }
     }
-    fn press(&mut self, target: u8) -> Vec<Vec<u8>> {
-        let path_options = self.child.as_mut().unwrap().press(target);
-        // for each path option, find our shortest route
-        let mut candidate_paths = Vec::new();
-        for child_path in path_options {
-            let candidate_path = self.path_to(&child_path);
-            candidate_paths.push(candidate_path);
+    fn press(&mut self, target: u8) -> Vec<FxHashMap<Vec<u8>, usize>> {
+        let child_frequencies = self.child.as_mut().unwrap().press(target);
+        let mut my_frequencies = Vec::new();
+
+        for freq_set in child_frequencies {
+            let mut local_freqs = FxHashMap::default();
+            for (moves, count) in freq_set {
+                assert_eq!(self.pointing_at, b'A');
+                let path = self.path_to(&moves);
+                for path_move in path.split_inclusive(|m| *m == b'A') {
+                    let entry = local_freqs.entry(path_move.to_vec()).or_insert(0);
+                    *entry = *entry + count;
+                }
+            }
+            my_frequencies.push(local_freqs);
         }
-        candidate_paths
+        my_frequencies
     }
 }
 
@@ -153,36 +180,43 @@ fn parse(input: &str) -> Vec<Code> {
     codes
 }
 
+fn run_robots(code: &Code, n: usize) -> i64 {
+    let numpad = Box::new(KeypadRobot::Number(NumberKeypadRobot::new()));
+
+    let mut robot = Box::new(KeypadRobot::Direction(DirectionKeypadRobot::new(0, Some(numpad))));
+    for i in 1..n {
+        let new_robot = Box::new(KeypadRobot::Direction(DirectionKeypadRobot::new(i, Some(robot))));
+        robot = new_robot
+    }
+
+    // let mut sol_freqs = Vec::new();
+    let mut sum = 0;
+    for button in &code.0 {
+        let paths = robot.press(*button).to_vec();
+        let best = paths
+            .iter()
+            .map(|bp| bp.iter().map(|(k, v)| k.len() * v).sum::<usize>())
+            .min()
+            .unwrap();
+        sum += best;
+    }
+    return sum as i64 * code.num_val();
+}
+
 #[aoc(day21, part1)]
 fn part1(input: &str) -> i64 {
     let codes = parse(input);
-    let mut sum = 0;
-    for code in &codes {
-        let numpad = NumberKeypadRobot::new();
-        let mut robot1 = DirectionKeypadRobot::new();
-        robot1.child = Some(numpad);
-
-        let mut robot2 = DirectionKeypadRobot::new();
-        robot2.child = Some(robot1);
-
-        let mut path = Vec::new();
-        for button in &code.0 {
-            let paths = robot2.press(*button);
-            path.push(paths);
-        }
-        let paths = path.clone().into_iter()
-            .multi_cartesian_product()
-            .map(|c| c.iter().flatten().map(|c| *c as char).join("")).collect_vec();
-        let best = paths.iter().map(|p| p.len()).min().unwrap() as i64;
-        let score = code.num_val() * best;
-        sum += score;
-    }
-    sum
+    codes.iter().map(|c| run_robots(c, 2)).sum::<i64>() as i64
 }
 
 #[aoc(day21, part2)]
 fn part2(input: &str) -> i64 {
-    todo!()
+    let codes = parse(input);
+    for i in 0..26 {
+        let res = codes.iter().map(|c| run_robots(c, i)).sum::<i64>() as i64;
+        println!("{i}: {res}");
+    }
+    0
 }
 
 #[cfg(test)]