aoc_2024/src/day21.rs

@@ -1,12 +1,14 @@
+use std::iter::once;
+
 use aoc_runner_derive::{aoc, aoc_generator};
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
-struct Path {
+struct Path<T> {
     steps: usize,
-    path: [Option<DirButton>; 6],
+    path: [Option<T>; 6],
 }
 
-impl Path {
+impl<T: Copy> Path<T> {
     fn new() -> Self {
         Self {
             steps: 0,
@@ -14,7 +16,7 @@ impl Path {
         }
     }
 
-    fn push(&mut self, button: DirButton) {
+    fn push(&mut self, button: T) {
         self.path[self.steps] = Some(button);
         self.steps += 1;
     }
@@ -24,36 +26,23 @@ impl Path {
             self.push(other.path[i].unwrap());
         }
     }
-
-    fn reversed(&self) -> Self {
-        let mut path = [None; 6];
-        (0..self.steps).for_each(|i| {
-            path[i] = match self.path[self.steps - i - 1] {
-                Some(DirButton::Up) => Some(DirButton::Down),
-                Some(DirButton::Down) => Some(DirButton::Up),
-                Some(DirButton::Left) => Some(DirButton::Right),
-                Some(DirButton::Right) => Some(DirButton::Left),
-                Some(DirButton::A) => Some(DirButton::A),
-                None => None,
-            };
-        });
-
-        Self {
-            steps: self.steps,
-            path,
-        }
-    }
 }
 
-impl IntoIterator for Path {
+impl<T> IntoIterator for Path<T> {
-    type Item = DirButton;
+    type Item = T;
-    type IntoIter = std::iter::Flatten<std::array::IntoIter<Option<DirButton>, 6>>;
+    type IntoIter = std::iter::Flatten<std::array::IntoIter<Option<T>, 6>>;
 
     fn into_iter(self) -> Self::IntoIter {
         self.path.into_iter().flatten()
     }
 }
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+enum Button {
+    Num(NumButton),
+    Dir(DirButton),
+}
+
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 enum DirButton {
     Up,
@@ -84,7 +73,7 @@ impl DirButton {
         }
     }
 
-    fn to(&self, other: &DirButton) -> [Option<Path>; 2] {
+    fn to(&self, other: &DirButton) -> [Option<Path<DirButton>>; 2] {
         let start = self.pos();
         let end = other.pos();
 
@@ -153,8 +142,7 @@ impl NumButton {
         }
     }
 
-    #[inline(always)]
+    fn to(&self, other: &NumButton) -> [Option<Path<DirButton>>; 2] {
-    fn to(&self, other: &NumButton) -> [Option<Path>; 2] {
         let start = self.pos();
         let end = other.pos();
 
@@ -223,67 +211,26 @@ fn parse(input: &str) -> Input {
 
 type DistanceMatrix = [[usize; 5]; 5];
 
-fn precompute_steps(steps: &mut [DistanceMatrix]) {
+fn calculate_steps(start: DirButton, end: DirButton, steps: &mut [DistanceMatrix]) -> usize {
-    if steps.len() == 1 {
+    if steps.is_empty() {
-        return;
+        return 1;
     }
 
-    precompute_steps(&mut steps[1..]);
+    if steps[0][start.index()][end.index()] == 0 {
+        let extension = start
+            .to(&end)
+            .into_iter()
+            .flatten()
+            .map(|path| shortest_dir_path(path, &mut steps[1..]))
+            .min()
+            .unwrap();
+        steps[0][start.index()][end.index()] = extension;
+    }
 
-    let buttons = [
+    steps[0][start.index()][end.index()]
-        DirButton::Up,
-        DirButton::Down,
-        DirButton::Left,
-        DirButton::Right,
-        DirButton::A,
-    ];
-
-    buttons.iter().enumerate().for_each(|(i, start)| {
-        buttons[i + 1..].iter().for_each(|end| {
-            let (forward, reversed) = start
-                .to(end)
-                .into_iter()
-                .flatten()
-                .map(|mut path| {
-                    path.steps -= 1;
-                    let mut reversed = path.reversed();
-                    let mut forward = path;
-                    reversed.push(DirButton::A);
-                    forward.push(DirButton::A);
-
-                    (
-                        forward
-                            .path
-                            .into_iter()
-                            .flatten()
-                            .fold((0, DirButton::A), |acc, next| {
-                                let extension = steps[1][acc.1.index()][next.index()];
-
-                                (acc.0 + extension, next)
-                            })
-                            .0,
-                        reversed
-                            .path
-                            .into_iter()
-                            .flatten()
-                            .fold((0, DirButton::A), |acc, next| {
-                                let extension = steps[1][acc.1.index()][next.index()];
-
-                                (acc.0 + extension, next)
-                            })
-                            .0,
-                    )
-                })
-                .fold((usize::MAX, usize::MAX), |acc, next| {
-                    (acc.0.min(next.0), acc.1.min(next.1))
-                });
-            steps[0][start.index()][end.index()] = forward;
-            steps[0][end.index()][start.index()] = reversed;
-        });
-    });
 }
 
-fn shortest_dir_path(path: Path, steps: &mut [DistanceMatrix]) -> usize {
+fn shortest_dir_path(path: Path<DirButton>, steps: &mut [DistanceMatrix]) -> usize {
     path.path
         .into_iter()
         .flatten()
@@ -291,7 +238,7 @@ fn shortest_dir_path(path: Path, steps: &mut [DistanceMatrix]) -> usize {
             let sum = acc.0;
             let pos = acc.1;
 
-            let extension = steps[0][pos.index()][next.index()];
+            let extension = calculate_steps(pos, next, steps);
 
             (sum + extension, next)
         })
@@ -318,9 +265,7 @@ fn shortest_path(path: Vec<NumButton>, steps: &mut [DistanceMatrix]) -> usize {
 
 #[aoc(day21, part1)]
 fn part1(input: &Input) -> usize {
-    let mut steps = vec![[[1; 5]; 5]; 3];
+    let mut steps = vec![[[0; 5]; 5]; 3 - 1];
-
-    precompute_steps(&mut steps);
 
     input
         .iter()
@@ -331,9 +276,7 @@ fn part1(input: &Input) -> usize {
 
 #[aoc(day21, part2)]
 fn part2(input: &Input) -> usize {
-    let mut steps = vec![[[1; 5]; 5]; 26];
+    let mut steps = vec![[[0; 5]; 5]; 26 - 1];
-
-    precompute_steps(&mut steps);
 
     input
         .iter()