From 2b379bba1d401a6127e3cabab2bb6e10e076efbc Mon Sep 17 00:00:00 2001
From: Jan-Bulthuis <git@bulthuis.dev>
Date: Sat, 21 Dec 2024 21:16:07 +0100
Subject: [PATCH] day21 fast

---
 aoc_2024/src/day21.rs | 343 +++++++++++++++++++-----------------------
 1 file changed, 156 insertions(+), 187 deletions(-)

diff --git a/aoc_2024/src/day21.rs b/aoc_2024/src/day21.rs
index 55b37f2..b8356d8 100644
--- a/aoc_2024/src/day21.rs
+++ b/aoc_2024/src/day21.rs
@@ -1,15 +1,14 @@
 use std::iter::once;
 
 use aoc_runner_derive::{aoc, aoc_generator};
-use hashbrown::{hash_map::Entry, HashMap};
 
-#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
-struct Path {
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+struct Path<T> {
     steps: usize,
-    path: [Option<Button>; 6],
+    path: [Option<T>; 6],
 }
 
-impl Path {
+impl<T: Copy> Path<T> {
     fn new() -> Self {
         Self {
             steps: 0,
@@ -17,7 +16,7 @@ impl Path {
         }
     }
 
-    fn push(&mut self, button: Button) {
+    fn push(&mut self, button: T) {
         self.path[self.steps] = Some(button);
         self.steps += 1;
     }
@@ -29,40 +28,23 @@ impl Path {
     }
 }
 
-impl IntoIterator for Path {
-    type Item = Button;
-    type IntoIter = std::iter::Flatten<std::array::IntoIter<Option<Button>, 6>>;
+impl<T> IntoIterator for Path<T> {
+    type Item = T;
+    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, Hash)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
 enum Button {
-    Directional(DirectionalButton),
-    Numeric(NumericButton),
+    Num(NumButton),
+    Dir(DirButton),
 }
 
-impl Button {
-    fn pos(&self) -> (isize, isize) {
-        match self {
-            Button::Directional(button) => button.pos(),
-            Button::Numeric(button) => button.pos(),
-        }
-    }
-
-    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),
-            _ => unimplemented!(),
-        }
-    }
-}
-
-#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
-enum DirectionalButton {
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+enum DirButton {
     Up,
     Down,
     Left,
@@ -70,18 +52,28 @@ enum DirectionalButton {
     A,
 }
 
-impl DirectionalButton {
-    fn pos(&self) -> (isize, isize) {
+impl DirButton {
+    fn pos(&self) -> (usize, usize) {
         match self {
-            DirectionalButton::Up => (1, 0),
-            DirectionalButton::Down => (1, 1),
-            DirectionalButton::Left => (0, 1),
-            DirectionalButton::Right => (2, 1),
-            DirectionalButton::A => (2, 0),
+            DirButton::Up => (1, 0),
+            DirButton::Down => (1, 1),
+            DirButton::Left => (0, 1),
+            DirButton::Right => (2, 1),
+            DirButton::A => (2, 0),
         }
     }
 
-    fn to(&self, other: &DirectionalButton) -> (Path, Option<Path>) {
+    fn index(&self) -> usize {
+        match self {
+            DirButton::Up => 0,
+            DirButton::Down => 1,
+            DirButton::Left => 2,
+            DirButton::Right => 3,
+            DirButton::A => 4,
+        }
+    }
+
+    fn to(&self, other: &DirButton) -> [Option<Path<DirButton>>; 2] {
         let start = self.pos();
         let end = other.pos();
 
@@ -89,59 +81,68 @@ impl DirectionalButton {
         let mut y = Path::new();
 
         if start.0 < end.0 {
-            (start.0..end.0).for_each(|_| x.push(Button::Directional(DirectionalButton::Right)));
+            (start.0..end.0).for_each(|_| x.push(DirButton::Right));
         } else {
-            (end.0..start.0).for_each(|_| x.push(Button::Directional(DirectionalButton::Left)));
+            (end.0..start.0).for_each(|_| x.push(DirButton::Left));
         };
         if start.1 < end.1 {
-            (start.1..end.1).for_each(|_| y.push(Button::Directional(DirectionalButton::Down)));
+            (start.1..end.1).for_each(|_| y.push(DirButton::Down));
         } else {
-            (end.1..start.1).for_each(|_| y.push(Button::Directional(DirectionalButton::Up)));
+            (end.1..start.1).for_each(|_| y.push(DirButton::Up));
         };
 
         if start.0 == 0 && end.1 == 0 {
             x.append(&y);
-            x.push(Button::Directional(DirectionalButton::A));
-            (x, None)
+            x.push(DirButton::A);
+            [Some(x), None]
         } else if end.0 == 0 && start.1 == 0 {
             y.append(&x);
-            y.push(Button::Directional(DirectionalButton::A));
-            (y, None)
+            y.push(DirButton::A);
+            [Some(y), None]
         } else {
+            let x_copy = x;
             x.append(&y);
-            x.push(Button::Directional(DirectionalButton::A));
-            y.append(&x);
-            y.push(Button::Directional(DirectionalButton::A));
-            (x, Some(y))
+            x.push(DirButton::A);
+            y.append(&x_copy);
+            y.push(DirButton::A);
+            [Some(x), Some(y)]
         }
     }
 }
 
-#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
-enum NumericButton {
-    Number(u8),
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+enum NumButton {
+    One,
+    Two,
+    Three,
+    Four,
+    Five,
+    Six,
+    Seven,
+    Eight,
+    Nine,
+    Zero,
     A,
 }
 
-impl NumericButton {
-    fn pos(&self) -> (isize, isize) {
+impl NumButton {
+    fn pos(&self) -> (usize, usize) {
         match self {
-            NumericButton::Number(1) => (0, 2),
-            NumericButton::Number(2) => (1, 2),
-            NumericButton::Number(3) => (2, 2),
-            NumericButton::Number(4) => (0, 1),
-            NumericButton::Number(5) => (1, 1),
-            NumericButton::Number(6) => (2, 1),
-            NumericButton::Number(7) => (0, 0),
-            NumericButton::Number(8) => (1, 0),
-            NumericButton::Number(9) => (2, 0),
-            NumericButton::Number(0) => (1, 3),
-            NumericButton::A => (2, 3),
-            _ => unreachable!(),
+            NumButton::One => (0, 2),
+            NumButton::Two => (1, 2),
+            NumButton::Three => (2, 2),
+            NumButton::Four => (0, 1),
+            NumButton::Five => (1, 1),
+            NumButton::Six => (2, 1),
+            NumButton::Seven => (0, 0),
+            NumButton::Eight => (1, 0),
+            NumButton::Nine => (2, 0),
+            NumButton::Zero => (1, 3),
+            NumButton::A => (2, 3),
         }
     }
 
-    fn to(&self, other: &NumericButton) -> (Path, Option<Path>) {
+    fn to(&self, other: &NumButton) -> [Option<Path<DirButton>>; 2] {
         let start = self.pos();
         let end = other.pos();
 
@@ -149,162 +150,138 @@ impl NumericButton {
         let mut y = Path::new();
 
         if start.0 < end.0 {
-            (start.0..end.0).for_each(|_| x.push(Button::Directional(DirectionalButton::Right)));
+            (start.0..end.0).for_each(|_| x.push(DirButton::Right));
         } else {
-            (end.0..start.0).for_each(|_| x.push(Button::Directional(DirectionalButton::Left)));
+            (end.0..start.0).for_each(|_| x.push(DirButton::Left));
         };
         if start.1 < end.1 {
-            (start.1..end.1).for_each(|_| y.push(Button::Directional(DirectionalButton::Down)));
+            (start.1..end.1).for_each(|_| y.push(DirButton::Down));
         } else {
-            (end.1..start.1).for_each(|_| y.push(Button::Directional(DirectionalButton::Up)));
+            (end.1..start.1).for_each(|_| y.push(DirButton::Up));
         };
 
         if start.0 == 0 && end.1 == 3 {
             x.append(&y);
-            x.push(Button::Directional(DirectionalButton::A));
-            (x, None)
+            x.push(DirButton::A);
+            [Some(x), None]
         } else if end.0 == 0 && start.1 == 3 {
             y.append(&x);
-            y.push(Button::Directional(DirectionalButton::A));
-            (y, None)
+            y.push(DirButton::A);
+            [Some(y), None]
         } else {
             let x_copy = x;
             x.append(&y);
-            x.push(Button::Directional(DirectionalButton::A));
+            x.push(DirButton::A);
             y.append(&x_copy);
-            y.push(Button::Directional(DirectionalButton::A));
-            (x, Some(y))
+            y.push(DirButton::A);
+            [Some(x), Some(y)]
         }
     }
 }
 
-type Input = Vec<Vec<Button>>;
+type Input = Vec<(usize, Vec<NumButton>)>;
 
 #[aoc_generator(day21)]
 fn parse(input: &str) -> Input {
     input
         .lines()
         .map(|line| {
-            line.chars()
-                .map(|c| match c {
-                    '0' => NumericButton::Number(0),
-                    '1' => NumericButton::Number(1),
-                    '2' => NumericButton::Number(2),
-                    '3' => NumericButton::Number(3),
-                    '4' => NumericButton::Number(4),
-                    '5' => NumericButton::Number(5),
-                    '6' => NumericButton::Number(6),
-                    '7' => NumericButton::Number(7),
-                    '8' => NumericButton::Number(8),
-                    '9' => NumericButton::Number(9),
-                    'A' => NumericButton::A,
-                    _ => unreachable!(),
-                })
-                .map(Button::Numeric)
-                .collect()
+            (
+                line[0..3].parse().unwrap(),
+                line.chars()
+                    .map(|c| match c {
+                        '1' => NumButton::One,
+                        '2' => NumButton::Two,
+                        '3' => NumButton::Three,
+                        '4' => NumButton::Four,
+                        '5' => NumButton::Five,
+                        '6' => NumButton::Six,
+                        '7' => NumButton::Seven,
+                        '8' => NumButton::Eight,
+                        '9' => NumButton::Nine,
+                        '0' => NumButton::Zero,
+                        'A' => NumButton::A,
+                        _ => unreachable!(),
+                    })
+                    .collect(),
+            )
         })
         .collect()
 }
 
-fn shortest_path(num_robots: usize, path: Path) -> usize {
-    type Memos = Vec<HashMap<Path, usize>>;
+type DistanceMatrix = [[usize; 5]; 5];
 
-    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
+fn calculate_steps(start: DirButton, end: DirButton, steps: &mut [DistanceMatrix]) -> usize {
+    if steps.is_empty() {
+        return 1;
     }
 
-    expand(num_robots, path, &mut memos)
+    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;
+    }
+
+    steps[0][start.index()][end.index()]
+}
+
+fn shortest_dir_path(path: Path<DirButton>, steps: &mut [DistanceMatrix]) -> usize {
+    path.path
+        .into_iter()
+        .flatten()
+        .fold((0, DirButton::A), |acc, next| {
+            let sum = acc.0;
+            let pos = acc.1;
+
+            let extension = calculate_steps(pos, next, steps);
+
+            (sum + extension, next)
+        })
+        .0
+}
+
+fn shortest_path(path: Vec<NumButton>, steps: &mut [DistanceMatrix]) -> usize {
+    path.into_iter()
+        .fold((0, NumButton::A), |acc, end| {
+            let sum = acc.0;
+            let start = acc.1;
+
+            let extension = start
+                .to(&end)
+                .into_iter()
+                .flatten()
+                .map(|path| shortest_dir_path(path, steps))
+                .min()
+                .unwrap();
+            (sum + extension, end)
+        })
+        .0
 }
 
 #[aoc(day21, part1)]
 fn part1(input: &Input) -> usize {
+    let mut steps = vec![[[0; 5]; 5]; 3 - 1];
+
     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(3, path);
-
-            numeric * path
-        })
+        .map(|(num, path)| num * shortest_path(path, &mut steps))
         .sum()
 }
 
 #[aoc(day21, part2)]
 fn part2(input: &Input) -> usize {
+    let mut steps = vec![[[0; 5]; 5]; 26 - 1];
+
     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
-        })
+        .map(|(num, path)| num * shortest_path(path, &mut steps))
         .sum()
 }
 
@@ -320,12 +297,4 @@ mod tests {
         assert_eq!(part1(&parse("980A")), 60 * 980);
         assert_eq!(part1(&parse("029A")), 68 * 29);
     }
-
-    // #[test]
-    // fn part2_example() {
-    //     assert_eq!(part2(&parse("<EXAMPLE>")), 0);
-    // }
 }
-
-// 89666 too low
-// 97150 wrong