Faster approach caching pairs of keypad buttons

Author: maneatingape

README.md

@@ -90,7 +90,7 @@ Performance is reasonable even on older hardware, for example a 2011 MacBook Pro
 | 18 | [RAM Run](https://adventofcode.com/2024/day/18) | [Source](src/year2024/day18.rs) | 42 |
 | 19 | [Linen Layout](https://adventofcode.com/2024/day/19) | [Source](src/year2024/day19.rs) | 118 |
 | 20 | [Race Condition](https://adventofcode.com/2024/day/20) | [Source](src/year2024/day20.rs) | 1354 |
-| 21 | [Keypad Conundrum](https://adventofcode.com/2024/day/21) | [Source](src/year2024/day21.rs) | 111 |
+| 21 | [Keypad Conundrum](https://adventofcode.com/2024/day/21) | [Source](src/year2024/day21.rs) | 20 |
 | 22 | [Monkey Market](https://adventofcode.com/2024/day/22) | [Source](src/year2024/day22.rs) | 1350 |
 | 23 | [LAN Party](https://adventofcode.com/2024/day/23) | [Source](src/year2024/day23.rs) | 43 |
 | 24 | [Crossed Wires](https://adventofcode.com/2024/day/24) | [Source](src/year2024/day24.rs) | 23 |

src/year2024/day21.rs

@@ -1,157 +1,138 @@
 //! # Keypad Conundrum
+//!
+//! Each key sequence always end in `A`. This means that we can consider each group of button
+//! presses between `A`s independently using a recursive approach with memoization to efficiently
+//! compute the minimum presses needed for any depth of chained robots.
 use crate::util::hash::*;
 use crate::util::parse::*;
 use crate::util::point::*;
+use std::iter::{once, repeat_n};
 
-type Cache = FastMap<(usize, usize), usize>;
+type Input = (Vec<(String, usize)>, Combinations);
+type Combinations = FastMap<(char, char), Vec<String>>;
+type Cache = FastMap<(char, char, usize), usize>;
 
-pub fn parse(input: &str) -> &str {
-    input
+/// Convert codes to pairs of the sequence itself with the numeric part.
+/// The pad combinations are the same between both parts so only need to be computed once.
+pub fn parse(input: &str) -> Input {
+    let pairs = input.lines().map(String::from).zip(input.iter_unsigned()).collect();
+    (pairs, pad_combinations())
 }
 
-pub fn part1(input: &str) -> usize {
+pub fn part1(input: &Input) -> usize {
     chain(input, 3)
 }
 
-pub fn part2(input: &str) -> usize {
+pub fn part2(input: &Input) -> usize {
     chain(input, 26)
 }
 
-fn chain(input: &str, limit: usize) -> usize {
+fn chain(input: &Input, depth: usize) -> usize {
+    let (pairs, combinations) = input;
     let cache = &mut FastMap::with_capacity(500);
-    input
-        .lines()
-        .map(str::as_bytes)
-        .zip(input.iter_unsigned::<usize>())
-        .map(|(code, numeric)| dfs(cache, code, 0, limit) * numeric)
-        .sum()
+    pairs.iter().map(|(code, numeric)| dfs(cache, combinations, code, depth) * numeric).sum()
 }
 
-fn dfs(cache: &mut Cache, slice: &[u8], depth: usize, limit: usize) -> usize {
-    if depth == limit {
-        return slice.len();
+fn dfs(cache: &mut Cache, combinations: &Combinations, code: &str, depth: usize) -> usize {
+    // Number of presses for the last keypad is just the length of the sequence.
+    if depth == 0 {
+        return code.len();
     }
 
-    let key = (to_usize(slice), depth);
-    if let Some(&previous) = cache.get(&key) {
-        return previous;
+    // All keypads start with `A`, either the initial position of the keypad or the trailing `A`
+    // from the previous sequence at this level.
+    let mut previous = 'A';
+    let mut result = 0;
+
+    for current in code.chars() {
+        // Check each pair of characters, memoizing results.
+        let key = (previous, current, depth);
+
+        result += cache.get(&key).copied().unwrap_or_else(|| {
+            // Each transition has either 1 or 2 possibilities.
+            // Pick the sequence that results in the minimum keypresses.
+            let presses = combinations[&(previous, current)]
+                .iter()
+                .map(|next| dfs(cache, combinations, next, depth - 1))
+                .min()
+                .unwrap();
+            cache.insert(key, presses);
+            presses
+        });
+
+        previous = current;
     }
 
-    let keypad = if depth == 0 { NUMERIC } else { DIRECTIONAL };
-    let mut shortest = usize::MAX;
-
-    for sequence in combinations(slice, &keypad) {
-        let mut presses = 0;
-
-        for chunk in sequence.split_inclusive(|&b| b == b'A') {
-            presses += dfs(cache, chunk, depth + 1, limit);
-        }
-
-        shortest = shortest.min(presses);
-    }
-
-    cache.insert(key, shortest);
-    shortest
+    result
 }
 
-fn combinations(current: &[u8], keypad: &Keypad) -> Vec<Vec<u8>> {
-    let mut next = Vec::new();
-    pad_dfs(&mut next, &mut Vec::with_capacity(16), keypad, current, 0, keypad.start);
-    next
+/// Compute keypresses needed for all possible transitions for both numeric and directional
+/// keypads. There are no distinct pairs shared between the keypads so they can use the same map
+/// without conflict.
+fn pad_combinations() -> Combinations {
+    let numeric_gap = Point::new(0, 3);
+    let numeric_keys = [
+        ('7', Point::new(0, 0)),
+        ('8', Point::new(1, 0)),
+        ('9', Point::new(2, 0)),
+        ('4', Point::new(0, 1)),
+        ('5', Point::new(1, 1)),
+        ('6', Point::new(2, 1)),
+        ('1', Point::new(0, 2)),
+        ('2', Point::new(1, 2)),
+        ('3', Point::new(2, 2)),
+        ('0', Point::new(1, 3)),
+        ('A', Point::new(2, 3)),
+    ];
+
+    let directional_gap = Point::new(0, 0);
+    let directional_keys = [
+        ('^', Point::new(1, 0)),
+        ('A', Point::new(2, 0)),
+        ('<', Point::new(0, 1)),
+        ('v', Point::new(1, 1)),
+        ('>', Point::new(2, 1)),
+    ];
+
+    let mut combinations = FastMap::new();
+    pad_routes(&mut combinations, &numeric_keys, numeric_gap);
+    pad_routes(&mut combinations, &directional_keys, directional_gap);
+    combinations
 }
 
-fn pad_dfs(
-    combinations: &mut Vec<Vec<u8>>,
-    path: &mut Vec<u8>,
-    keypad: &Keypad,
-    sequence: &[u8],
-    depth: usize,
-    from: Point,
-) {
-    // Success
-    if depth == sequence.len() {
-        combinations.push(path.clone());
-        return;
-    }
-
-    // Failure
-    if from == keypad.gap {
-        return;
-    }
-
-    let to = keypad.lookup[sequence[depth] as usize];
-
-    if from == to {
-        // Push button.
-        path.push(b'A');
-        pad_dfs(combinations, path, keypad, sequence, depth + 1, from);
-        path.pop();
-    } else {
-        // Move towards button.
-        let mut step = |next: u8, direction: Point| {
-            path.push(next);
-            pad_dfs(combinations, path, keypad, sequence, depth, from + direction);
-            path.pop();
-        };
-
-        if to.x < from.x {
-            step(b'<', LEFT);
-        }
-        if to.x > from.x {
-            step(b'>', RIGHT);
-        }
-        if to.y < from.y {
-            step(b'^', UP);
-        }
-        if to.y > from.y {
-            step(b'v', DOWN);
+/// Each route between two keys has 2 possibilites, horizontal first or vertical first.
+/// We skip any route that would cross the gap and also avoid adding the same route twice
+/// when a key is in a straight line (e.g. directly above/below or left/right). For example:
+///
+/// * `7 => A` is only `>>vvv`.
+/// * `1 => 5` is `^>` and `>^`.
+///
+/// We don't consider routes that change direction more than once as these are always longer,
+/// for example `5 => A` ignores the path `v>v`.
+fn pad_routes(combinations: &mut Combinations, pad: &[(char, Point)], gap: Point) {
+    for &(first, from) in pad {
+        for &(second, to) in pad {
+            let horizontal = || {
+                let element = if from.x < to.x { '>' } else { '<' };
+                let count = from.x.abs_diff(to.x) as usize;
+                repeat_n(element, count)
+            };
+
+            let vertical = || {
+                let element = if from.y < to.y { 'v' } else { '^' };
+                let count = from.y.abs_diff(to.y) as usize;
+                repeat_n(element, count)
+            };
+
+            if Point::new(from.x, to.y) != gap {
+                let path = vertical().chain(horizontal()).chain(once('A')).collect();
+                combinations.entry((first, second)).or_default().push(path);
+            }
+
+            if from.x != to.x && from.y != to.y && Point::new(to.x, from.y) != gap {
+                let path = horizontal().chain(vertical()).chain(once('A')).collect();
+                combinations.entry((first, second)).or_default().push(path);
+            }
         }
     }
 }
-
-struct Keypad {
-    start: Point,
-    gap: Point,
-    lookup: [Point; 128],
-}
-
-const NUMERIC: Keypad = {
-    let start = Point::new(2, 3);
-    let gap = Point::new(0, 3);
-    let mut lookup = [ORIGIN; 128];
-
-    lookup[b'7' as usize] = Point::new(0, 0);
-    lookup[b'8' as usize] = Point::new(1, 0);
-    lookup[b'9' as usize] = Point::new(2, 0);
-    lookup[b'4' as usize] = Point::new(0, 1);
-    lookup[b'5' as usize] = Point::new(1, 1);
-    lookup[b'6' as usize] = Point::new(2, 1);
-    lookup[b'1' as usize] = Point::new(0, 2);
-    lookup[b'2' as usize] = Point::new(1, 2);
-    lookup[b'3' as usize] = Point::new(2, 2);
-    lookup[b'0' as usize] = Point::new(1, 3);
-    lookup[b'A' as usize] = Point::new(2, 3);
-
-    Keypad { start, gap, lookup }
-};
-
-const DIRECTIONAL: Keypad = {
-    let start = Point::new(2, 0);
-    let gap = Point::new(0, 0);
-    let mut lookup = [ORIGIN; 128];
-
-    lookup[b'^' as usize] = Point::new(1, 0);
-    lookup[b'A' as usize] = Point::new(2, 0);
-    lookup[b'<' as usize] = Point::new(0, 1);
-    lookup[b'v' as usize] = Point::new(1, 1);
-    lookup[b'>' as usize] = Point::new(2, 1);
-
-    Keypad { start, gap, lookup }
-};
-
-// Max slice length is 5 so value is unique.
-fn to_usize(slice: &[u8]) -> usize {
-    let mut array = [0; 8];
-    array[0..slice.len()].copy_from_slice(slice);
-    usize::from_ne_bytes(array)
-}

tests/year2024/day21.rs

@@ -10,11 +10,11 @@ const EXAMPLE: &str = "\
 #[test]
 fn part1_test() {
     let input = parse(EXAMPLE);
-    assert_eq!(part1(input), 126384);
+    assert_eq!(part1(&input), 126384);
 }
 
 #[test]
 fn part2_test() {
     let input = parse(EXAMPLE);
-    assert_eq!(part2(input), 154115708116294);
+    assert_eq!(part2(&input), 154115708116294);
 }