day21 working but not fast enough for part 2

aoc_2024/src/day21.rs

@@ -1,6 +1,44 @@
-use aoc_runner_derive::{aoc, aoc_generator};
+use std::iter::once;
 
-#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+use aoc_runner_derive::{aoc, aoc_generator};
+use hashbrown::{hash_map::Entry, HashMap};
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+struct Path {
+    steps: usize,
+    path: [Option<Button>; 6],
+}
+
+impl Path {
+    fn new() -> Self {
+        Self {
+            steps: 0,
+            path: [None, None, None, None, None, None],
+        }
+    }
+
+    fn push(&mut self, button: Button) {
+        self.path[self.steps] = Some(button);
+        self.steps += 1;
+    }
+
+    fn append(&mut self, other: &Self) {
+        for i in 0..other.steps {
+            self.push(other.path[i].unwrap());
+        }
+    }
+}
+
+impl IntoIterator for Path {
+    type Item = Button;
+    type IntoIter = std::iter::Flatten<std::array::IntoIter<Option<Button>, 6>>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.path.into_iter().flatten()
+    }
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
 enum Button {
     Directional(DirectionalButton),
     Numeric(NumericButton),
@@ -14,7 +52,7 @@ impl Button {
         }
     }
 
-    fn to(&self, other: &Button) -> Vec<DirectionalButton> {
+    fn to(&self, other: &Button) -> (Path, Option<Path>) {
         match (self, other) {
             (Button::Directional(a), Button::Directional(b)) => a.to(b),
             (Button::Numeric(a), Button::Numeric(b)) => a.to(b),
@@ -23,7 +61,7 @@ impl Button {
     }
 }
 
-#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
 enum DirectionalButton {
     Up,
     Down,
@@ -43,30 +81,43 @@ impl DirectionalButton {
         }
     }
 
-    fn to(&self, other: &DirectionalButton) -> Vec<DirectionalButton> {
+    fn to(&self, other: &DirectionalButton) -> (Path, Option<Path>) {
         let start = self.pos();
         let end = other.pos();
-        let mut x = if start.0 < end.0 {
+
-            vec![DirectionalButton::Right; (end.0 - start.0) as usize]
+        let mut x = Path::new();
+        let mut y = Path::new();
+
+        if start.0 < end.0 {
+            (start.0..end.0).for_each(|_| x.push(Button::Directional(DirectionalButton::Right)));
         } else {
-            vec![DirectionalButton::Left; (start.0 - end.0) as usize]
+            (end.0..start.0).for_each(|_| x.push(Button::Directional(DirectionalButton::Left)));
         };
-        let mut y = if start.1 < end.1 {
+        if start.1 < end.1 {
-            vec![DirectionalButton::Down; (end.1 - start.1) as usize]
+            (start.1..end.1).for_each(|_| y.push(Button::Directional(DirectionalButton::Down)));
         } else {
-            vec![DirectionalButton::Up; (start.1 - end.1) as usize]
+            (end.1..start.1).for_each(|_| y.push(Button::Directional(DirectionalButton::Up)));
         };
-        if end.0 == 0 {
+
-            y.append(&mut x);
+        if start.0 == 0 && end.1 == 0 {
-            y
+            x.append(&y);
+            x.push(Button::Directional(DirectionalButton::A));
+            (x, None)
+        } else if end.0 == 0 && start.1 == 0 {
+            y.append(&x);
+            y.push(Button::Directional(DirectionalButton::A));
+            (y, None)
         } else {
-            x.append(&mut y);
+            x.append(&y);
-            x
+            x.push(Button::Directional(DirectionalButton::A));
+            y.append(&x);
+            y.push(Button::Directional(DirectionalButton::A));
+            (x, Some(y))
         }
     }
 }
 
-#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
 enum NumericButton {
     Number(u8),
     A,
@@ -90,25 +141,39 @@ impl NumericButton {
         }
     }
 
-    fn to(&self, other: &NumericButton) -> Vec<DirectionalButton> {
+    fn to(&self, other: &NumericButton) -> (Path, Option<Path>) {
         let start = self.pos();
         let end = other.pos();
-        let mut x = if start.0 < end.0 {
+
-            vec![DirectionalButton::Right; (end.0 - start.0) as usize]
+        let mut x = Path::new();
+        let mut y = Path::new();
+
+        if start.0 < end.0 {
+            (start.0..end.0).for_each(|_| x.push(Button::Directional(DirectionalButton::Right)));
         } else {
-            vec![DirectionalButton::Left; (start.0 - end.0) as usize]
+            (end.0..start.0).for_each(|_| x.push(Button::Directional(DirectionalButton::Left)));
         };
-        let mut y = if start.1 < end.1 {
+        if start.1 < end.1 {
-            vec![DirectionalButton::Down; (end.1 - start.1) as usize]
+            (start.1..end.1).for_each(|_| y.push(Button::Directional(DirectionalButton::Down)));
         } else {
-            vec![DirectionalButton::Up; (start.1 - end.1) as usize]
+            (end.1..start.1).for_each(|_| y.push(Button::Directional(DirectionalButton::Up)));
         };
-        if end.1 != 3 {
+
-            y.append(&mut x);
+        if start.0 == 0 && end.1 == 3 {
-            y
+            x.append(&y);
+            x.push(Button::Directional(DirectionalButton::A));
+            (x, None)
+        } else if end.0 == 0 && start.1 == 3 {
+            y.append(&x);
+            y.push(Button::Directional(DirectionalButton::A));
+            (y, None)
         } else {
-            x.append(&mut y);
+            let x_copy = x;
-            x
+            x.append(&y);
+            x.push(Button::Directional(DirectionalButton::A));
+            y.append(&x_copy);
+            y.push(Button::Directional(DirectionalButton::A));
+            (x, Some(y))
         }
     }
 }
@@ -141,45 +206,56 @@ fn parse(input: &str) -> Input {
         .collect()
 }
 
+fn shortest_path(num_robots: usize, path: Path) -> usize {
+    type Memos = Vec<HashMap<Path, usize>>;
+
+    let mut memos = vec![HashMap::new(); num_robots];
+
+    fn expand(robots: usize, path: Path, memos: &mut Memos) -> usize {
+        if robots == 0 {
+            return path.steps;
+        }
+
+        let pos = match path.path[0].unwrap() {
+            Button::Numeric(_) => Button::Numeric(NumericButton::A),
+            Button::Directional(_) => Button::Directional(DirectionalButton::A),
+        };
+
+        path.into_iter()
+            .fold((pos, 0), |acc, next| {
+                let pos = acc.0;
+                let steps = acc.1;
+
+                let (first, second) = pos.to(&next);
+                let expansion = once(Some(first))
+                    .chain(once(second))
+                    .flatten()
+                    .min_by_key(|path| {
+                        if memos[robots - 1].contains_key(path) {
+                            memos[robots - 1][path]
+                        } else {
+                            let expanded = expand(robots - 1, *path, memos);
+                            memos[robots - 1].insert(*path, expanded);
+                            expanded
+                        }
+                    })
+                    .map(|path| expand(robots - 1, path, memos))
+                    .unwrap();
+
+                (next, steps + expansion)
+            })
+            .1
+    }
+
+    expand(num_robots, path, &mut memos)
+}
+
 #[aoc(day21, part1)]
 fn part1(input: &Input) -> usize {
     input
         .iter()
         .cloned()
         .map(|instruction| {
-            let mut instructions = vec![];
-            let mut stack = vec![];
-            let num_robots = 3;
-            let mut positions = vec![Button::Numeric(NumericButton::A)];
-            for _ in 1..num_robots {
-                positions.push(Button::Directional(DirectionalButton::A));
-            }
-            for button in instruction.iter().rev() {
-                stack.push((0, *button));
-            }
-
-            while let Some((robot, button)) = stack.pop() {
-                if robot == num_robots {
-                    instructions.push(button);
-                    continue;
-                }
-
-                println!(
-                    "move {} from {:?} to {:?} = {:?}",
-                    robot,
-                    positions[robot].pos(),
-                    button.pos(),
-                    button
-                );
-
-                let instructions = positions[robot].to(&button);
-                stack.push((robot + 1, Button::Directional(DirectionalButton::A)));
-                for instruction in instructions.iter().rev() {
-                    stack.push((robot + 1, Button::Directional(*instruction)));
-                }
-                positions[robot] = button;
-            }
-
             let numeric = instruction
                 .iter()
                 .filter_map(|button| match button {
@@ -193,30 +269,44 @@ fn part1(input: &Input) -> usize {
                 .parse::<usize>()
                 .unwrap();
 
-            dbg!(instructions.len());
+            let mut path = Path::new();
-            println!(
+            instruction.iter().for_each(|step| path.push(*step));
-                "{}",
+
-                instructions
+            let path = shortest_path(3, path);
-                    .iter()
+
-                    .map(|instruction| match instruction {
+            numeric * path
-                        Button::Directional(DirectionalButton::Up) => '^',
-                        Button::Directional(DirectionalButton::Down) => 'v',
-                        Button::Directional(DirectionalButton::Left) => '<',
-                        Button::Directional(DirectionalButton::Right) => '>',
-                        Button::Directional(DirectionalButton::A) => 'A',
-                        _ => unreachable!(),
-                    })
-                    .collect::<String>()
-            );
-            numeric * instructions.len()
         })
         .sum()
 }
 
-// #[aoc(day21, part2)]
+#[aoc(day21, part2)]
-// fn part2(input: &Input) -> usize {
+fn part2(input: &Input) -> usize {
-//     todo!()
+    input
-// }
+        .iter()
+        .cloned()
+        .map(|instruction| {
+            let numeric = instruction
+                .iter()
+                .filter_map(|button| match button {
+                    Button::Numeric(NumericButton::Number(numeric)) => Some(numeric),
+                    _ => None,
+                })
+                .fold(String::new(), |mut acc, num| {
+                    acc.push_str(&num.to_string());
+                    acc
+                })
+                .parse::<usize>()
+                .unwrap();
+
+            let mut path = Path::new();
+            instruction.iter().for_each(|step| path.push(*step));
+
+            let path = shortest_path(8, path);
+
+            numeric * path
+        })
+        .sum()
+}
 
 #[cfg(test)]
 mod tests {
@@ -224,7 +314,7 @@ mod tests {
 
     #[test]
     fn part1_example() {
-        // assert_eq!(part1(&parse("179A")), 68 * 179);
+        assert_eq!(part1(&parse("179A")), 68 * 179);
         assert_eq!(part1(&parse("379A")), 64 * 379);
         assert_eq!(part1(&parse("456A")), 64 * 456);
         assert_eq!(part1(&parse("980A")), 60 * 980);