Rust/2024/21: add solution

Author: Defelo

README.md

@@ -0,0 +1,172 @@
+#![feature(test)]
+
+use std::iter;
+
+use aoc::matrix;
+use itertools::Itertools;
+use rustc_hash::FxHashSet;
+
+type Input = Vec<Vec<u8>>;
+
+fn setup(input: &str) -> Input {
+    input
+        .trim()
+        .lines()
+        .map(|l| {
+            l.bytes()
+                .map(|b| match b {
+                    b'0'..=b'9' => b - b'0',
+                    b'A' => 10,
+                    _ => panic!(),
+                })
+                .collect()
+        })
+        .collect()
+}
+
+type Sequence = Vec<u8>;
+
+/// For each keypad button pair compute all optimal sequences to activate the
+/// second button, starting at the first one.
+fn compute_sequences<const N: usize>(keypad: [(u8, u8); N]) -> [[Vec<Sequence>; N]; N] {
+    keypad.map(|(i1, j1)| {
+        keypad.map(|(i2, j2)| {
+            iter::empty()
+                .chain(iter::repeat_n(0, i1.saturating_sub(i2) as _))
+                .chain(iter::repeat_n(1, j1.saturating_sub(j2) as _))
+                .chain(iter::repeat_n(2, i2.saturating_sub(i1) as _))
+                .chain(iter::repeat_n(3, j2.saturating_sub(j1) as _))
+                .permutations((i1.abs_diff(i2) + j1.abs_diff(j2)) as _)
+                .filter(|seq| {
+                    let (mut i, mut j) = (i1, j1);
+                    seq.iter().all(|&d| {
+                        match d {
+                            0 => i -= 1,
+                            1 => j -= 1,
+                            2 => i += 1,
+                            3 => j += 1,
+                            _ => panic!(),
+                        }
+                        keypad.contains(&(i, j))
+                    })
+                })
+                .collect::<FxHashSet<_>>()
+                .into_iter()
+                .map(|mut s| {
+                    s.push(DIR_A);
+                    s
+                })
+                .collect()
+        })
+    })
+}
+
+const NUM_A: u8 = 10;
+const DIR_A: u8 = 4;
+
+type NumSeqs = [[Vec<Sequence>; 11]; 11];
+type DirSeqs = [[Vec<Sequence>; 5]; 5];
+
+const NUM_KEYPAD: [(u8, u8); 11] = [
+    (3, 1), // 0
+    (2, 0), // 1
+    (2, 1), // 2
+    (2, 2), // 3
+    (1, 0), // 4
+    (1, 1), // 5
+    (1, 2), // 6
+    (0, 0), // 7
+    (0, 1), // 8
+    (0, 2), // 9
+    (3, 2), // A
+];
+
+const DIR_KEYPAD: [(u8, u8); 5] = [
+    (0, 1), // ^
+    (1, 0), // <
+    (1, 1), // v
+    (1, 2), // >
+    (0, 2), // A
+];
+
+fn steps(init: u8, seq: &[u8]) -> impl Iterator<Item = (usize, usize)> + use<'_> {
+    iter::once(init as _)
+        .chain(seq.iter().map(|&x| x as _))
+        .tuple_windows()
+}
+
+type DirMatrix = [[usize; 5]; 5];
+
+/// For each button pair on the directional keypad closest to the door compute
+/// the minimum number of button presses needed on the manually controlled
+/// directional keypad to activate the second button starting at the first one.
+fn compute_dir_matrix<const N: usize>(dir_seqs: &DirSeqs) -> DirMatrix {
+    (1..N).fold(
+        // if there is only one robot controlled directional keypad, use one of the pre-computed
+        // sequences on the human controlled keypad. doesn't matter which one, they should all have
+        // the same (optimal) length.
+        matrix(|x, y| dir_seqs[x][y][0].len()),
+        // if there are more robot controlled directional keypads, compute a new matrix based on
+        // the previous one: try all sequences from x to y on the directional keypad, sum up all
+        // the required button presses for each sequence and find the minimum of those.
+        |prev, _| {
+            matrix(|x, y| {
+                dir_seqs[x][y]
+                    .iter()
+                    .map(|seq| steps(DIR_A, seq).map(|(x, y)| prev[x][y]).sum())
+                    .min()
+                    .unwrap()
+            })
+        },
+    )
+}
+
+type NumMatrix = [[usize; 11]; 11];
+
+/// Similarly to [`compute_dir_matrix`], for each button pair on the numeric
+/// keypad at the door compute the minimum number of button presses needed on
+/// the manually controlled directional keypad to activate the second button
+/// starting at the first one.
+fn compute_num_matrix(num_seqs: &NumSeqs, dir_matrix: &DirMatrix) -> NumMatrix {
+    matrix(|x, y| {
+        // try all sequences from x to y on the numeric keypad, sum up all the required
+        // button presses for each sequence and find the minimum of those.
+        num_seqs[x][y]
+            .iter()
+            .map(|seq| steps(DIR_A, seq).map(|(x, y)| dir_matrix[x][y]).sum())
+            .min()
+            .unwrap()
+    })
+}
+
+fn min_len(seq: &[u8], num_matrix: &NumMatrix) -> usize {
+    steps(NUM_A, seq).map(|(x, y)| num_matrix[x][y]).sum()
+}
+
+fn numeric_part(seq: &[u8]) -> usize {
+    seq.iter()
+        .filter(|&&b| b < 10)
+        .fold(0, |acc, &b| acc * 10 + b as usize)
+}
+
+fn solve<const N: usize>(input: &Input) -> usize {
+    let num_seqs = compute_sequences(NUM_KEYPAD);
+    let dir_seqs = compute_sequences(DIR_KEYPAD);
+    let dir_matrix = compute_dir_matrix::<N>(&dir_seqs);
+    let num_matrix = compute_num_matrix(&num_seqs, &dir_matrix);
+
+    input
+        .iter()
+        .map(|seq| min_len(seq, &num_matrix) * numeric_part(seq))
+        .sum()
+}
+
+fn part1(input: &Input) -> usize {
+    solve::<2>(input)
+}
+
+fn part2(input: &Input) -> usize {
+    solve::<25>(input)
+}
+
+aoc::main!(2024, 21, ex: 1);

Rust/2024/21.rs

@@ -19,7 +19,7 @@ aoc::year! {
     "18.rs",
     "19.rs",
     "20.rs",
-    // "21.rs",
+    "21.rs",
     // "22.rs",
     // "23.rs",
     // "24.rs",

Rust/2024/main.rs

@@ -420,3 +420,6 @@ path = "2024/19.rs"
 [[bin]]
 name = "2024_20"
 path = "2024/20.rs"
+[[bin]]
+name = "2024_21"
+path = "2024/21.rs"

Rust/Cargo.toml

@@ -15,6 +15,12 @@ pub mod tuples;
 
 extern crate test;
 
+pub fn matrix<const N: usize, const M: usize>(
+    mut f: impl FnMut(usize, usize) -> usize,
+) -> [[usize; M]; N] {
+    std::array::from_fn(|x| std::array::from_fn(|y| f(x, y)))
+}
+
 #[macro_export]
 macro_rules! main {
     ($year:literal, $day:tt $(, ex: $($example:literal $([$expart:ident])?),*)?) => {

Rust/lib/lib.rs

@@ -0,0 +1,5 @@
+029A
+980A
+179A
+456A
+379A

examples/2024/21/1

@@ -0,0 +1 @@
+126384

examples/2024/21/1.1

@@ -0,0 +1 @@
+154115708116294