Skip to content

Commit 435bc00

Browse files
ictrobotictrobot
committed
Compute 2024 day 21 cost matrices using const fns at compile time
1 parent cac0e50 commit 435bc00

File tree

1 file changed

+130
-74
lines changed

1 file changed

+130
-74
lines changed

‎crates/year2024/src/day21.rs‎

Lines changed: 130 additions & 74 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,4 @@
1-
use std::cmp::Reverse;
2-
use std::collections::BinaryHeap;
1+
use utils::point::Point2D;
32
use utils::prelude::*;
43

54
/// Counting recursive keypad presses.
@@ -24,6 +23,9 @@ enum DirectionalKeypad {
2423
Left = 1, Down = 2, Right = 3,
2524
}
2625

26+
static PART1_MATRIX: [[u64; 11]; 11] = num_matrix(2);
27+
static PART2_MATRIX: [[u64; 11]; 11] = num_matrix(25);
28+
2729
impl Day21 {
2830
pub fn new(input: &str, _: InputType) -> Result<Self, InputError> {
2931
Ok(Self {
@@ -35,80 +37,137 @@ impl Day21 {
3537

3638
#[must_use]
3739
pub fn part1(&self) -> u64 {
38-
self.complexity(2)
40+
self.complexity(&PART1_MATRIX)
3941
}
4042

4143
#[must_use]
4244
pub fn part2(&self) -> u64 {
43-
self.complexity(25)
45+
self.complexity(&PART2_MATRIX)
4446
}
4547

46-
fn complexity(&self, robots: u32) -> u64 {
47-
let mut dir_matrix = [[1; 5]; 5];
48-
for _ in 0..robots {
49-
dir_matrix = Self::dir_cost(dir_matrix);
50-
}
51-
52-
let num_matrix = Self::num_cost(dir_matrix);
53-
48+
fn complexity(&self, matrix: &[[u64; 11]; 11]) -> u64 {
5449
self.codes
5550
.iter()
5651
.map(|&code| {
57-
let digits = [code / 100, (code % 100) / 10, code % 10];
58-
let length = num_matrix[NumericKeypad::Activate as usize][digits[0] as usize]
59-
+ num_matrix[digits[0] as usize][digits[1] as usize]
60-
+ num_matrix[digits[1] as usize][digits[2] as usize]
61-
+ num_matrix[digits[2] as usize][NumericKeypad::Activate as usize];
52+
let digits = [(code % 1000) / 100, (code % 100) / 10, code % 10];
53+
let length = matrix[NumericKeypad::Activate as usize][digits[0] as usize]
54+
+ matrix[digits[0] as usize][digits[1] as usize]
55+
+ matrix[digits[1] as usize][digits[2] as usize]
56+
+ matrix[digits[2] as usize][NumericKeypad::Activate as usize];
6257
length * code as u64
6358
})
6459
.sum()
6560
}
61+
}
6662

67-
fn dir_cost(
68-
dir_matrix: [[u64; DirectionalKeypad::LEN]; DirectionalKeypad::LEN],
69-
) -> [[u64; DirectionalKeypad::LEN]; DirectionalKeypad::LEN] {
70-
let mut result = [[u64::MAX; DirectionalKeypad::LEN]; DirectionalKeypad::LEN];
71-
let mut queue = BinaryHeap::new();
72-
for start in DirectionalKeypad::ALL {
73-
queue.push(Reverse((0, start, DirectionalKeypad::Activate)));
74-
result[start as usize][start as usize] = 1;
75-
while let Some(Reverse((cost, current, parent))) = queue.pop() {
76-
for &(next_parent, next) in current.neighbours() {
77-
let next_cost = cost + dir_matrix[parent as usize][next_parent as usize];
78-
let activate_cost = next_cost
79-
+ dir_matrix[next_parent as usize][DirectionalKeypad::Activate as usize];
80-
if result[start as usize][next as usize] > activate_cost {
81-
result[start as usize][next as usize] = activate_cost;
82-
queue.push(Reverse((next_cost, next, next_parent)));
83-
}
84-
}
63+
const fn num_matrix(robots: u32) -> [[u64; NumericKeypad::LEN]; NumericKeypad::LEN] {
64+
let mut dir_matrix = [[1; 5]; 5];
65+
let mut i = 0;
66+
while i < robots {
67+
dir_matrix = DirectionalKeypad::cost_matrix(dir_matrix);
68+
i += 1;
69+
}
70+
NumericKeypad::cost_matrix(dir_matrix)
71+
}
72+
73+
// Use a macro for common functions as trait functions cannot be marked as const
74+
macro_rules! cost_matrix_functions {
75+
() => {
76+
const fn cost_matrix(
77+
dir_matrix: [[u64; DirectionalKeypad::LEN]; DirectionalKeypad::LEN],
78+
) -> [[u64; Self::LEN]; Self::LEN] {
79+
let mut result = [[u64::MAX; Self::LEN]; Self::LEN];
80+
let mut i = 0;
81+
while i < Self::LEN {
82+
result[i][i] = 1;
83+
Self::visit(
84+
0,
85+
Self::ALL[i],
86+
DirectionalKeypad::Activate,
87+
Self::ALL[i],
88+
dir_matrix,
89+
&mut result,
90+
);
91+
i += 1;
8592
}
93+
result
8694
}
87-
result
88-
}
8995

90-
fn num_cost(
91-
dir_matrix: [[u64; DirectionalKeypad::LEN]; DirectionalKeypad::LEN],
92-
) -> [[u64; NumericKeypad::LEN]; NumericKeypad::LEN] {
93-
let mut result = [[u64::MAX; NumericKeypad::LEN]; NumericKeypad::LEN];
94-
let mut queue = BinaryHeap::new();
95-
for start in NumericKeypad::ALL {
96-
queue.push(Reverse((0, start, DirectionalKeypad::Activate)));
97-
result[start as usize][start as usize] = 1;
98-
while let Some(Reverse((cost, current, parent))) = queue.pop() {
99-
for &(next_parent, next) in current.neighbours() {
100-
let next_cost = cost + dir_matrix[parent as usize][next_parent as usize];
101-
let activate_cost = next_cost
102-
+ dir_matrix[next_parent as usize][DirectionalKeypad::Activate as usize];
103-
if result[start as usize][next as usize] > activate_cost {
104-
result[start as usize][next as usize] = activate_cost;
105-
queue.push(Reverse((next_cost, next, next_parent)));
106-
}
96+
const fn visit(
97+
cost: u64,
98+
current: Self,
99+
parent: DirectionalKeypad,
100+
start: Self,
101+
dir_matrix: [[u64; DirectionalKeypad::LEN]; DirectionalKeypad::LEN],
102+
result: &mut [[u64; Self::LEN]; Self::LEN],
103+
) {
104+
let cost_with_activate =
105+
cost + dir_matrix[parent as usize][DirectionalKeypad::Activate as usize];
106+
if cost_with_activate < result[start as usize][current as usize] {
107+
result[start as usize][current as usize] = cost_with_activate;
108+
}
109+
110+
let start_coords = start.coords();
111+
let current_coords = current.coords();
112+
let current_distance = current_coords.x.abs_diff(start_coords.x)
113+
+ current_coords.y.abs_diff(start_coords.y);
114+
115+
let neighbours = current.neighbours();
116+
let mut i = 0;
117+
while i < neighbours.len() {
118+
let (next_parent, next) = neighbours[i];
119+
let next_point = next.coords();
120+
let next_distance =
121+
next_point.x.abs_diff(start_coords.x) + next_point.y.abs_diff(start_coords.y);
122+
if next_distance > current_distance {
123+
Self::visit(
124+
cost + dir_matrix[parent as usize][next_parent as usize],
125+
next,
126+
next_parent,
127+
start,
128+
dir_matrix,
129+
result,
130+
);
107131
}
132+
i += 1;
108133
}
109134
}
110-
result
135+
};
136+
}
137+
138+
impl DirectionalKeypad {
139+
const ALL: [Self; 5] = [
140+
DirectionalKeypad::Up,
141+
DirectionalKeypad::Activate,
142+
DirectionalKeypad::Left,
143+
DirectionalKeypad::Down,
144+
DirectionalKeypad::Right,
145+
];
146+
const LEN: usize = 5;
147+
148+
const fn neighbours(self) -> &'static [(DirectionalKeypad, Self)] {
149+
use DirectionalKeypad::*;
150+
match self {
151+
Up => &[(Right, Activate), (Down, Down)],
152+
Activate => &[(Left, Up), (Down, Right)],
153+
Left => &[(Right, Down)],
154+
Down => &[(Up, Up), (Left, Left), (Right, Right)],
155+
Right => &[(Up, Activate), (Left, Down)],
156+
}
157+
}
158+
159+
const fn coords(self) -> Point2D<u32> {
160+
use DirectionalKeypad::*;
161+
match self {
162+
Up => Point2D::new(1, 0),
163+
Activate => Point2D::new(2, 0),
164+
Left => Point2D::new(0, 1),
165+
Down => Point2D::new(1, 1),
166+
Right => Point2D::new(2, 1),
167+
}
111168
}
169+
170+
cost_matrix_functions!();
112171
}
113172

114173
impl NumericKeypad {
@@ -127,7 +186,7 @@ impl NumericKeypad {
127186
];
128187
const LEN: usize = 11;
129188

130-
fn neighbours(self) -> &'static [(DirectionalKeypad, Self)] {
189+
const fn neighbours(self) -> &'static [(DirectionalKeypad, Self)] {
131190
use DirectionalKeypad::{Down, Left, Right, Up};
132191
use NumericKeypad::*;
133192
match self {
@@ -144,28 +203,25 @@ impl NumericKeypad {
144203
Activate => &[(Up, Key3), (Left, Key0)],
145204
}
146205
}
147-
}
148-
149-
impl DirectionalKeypad {
150-
const ALL: [Self; 5] = [
151-
DirectionalKeypad::Up,
152-
DirectionalKeypad::Activate,
153-
DirectionalKeypad::Left,
154-
DirectionalKeypad::Down,
155-
DirectionalKeypad::Right,
156-
];
157-
const LEN: usize = 5;
158206

159-
fn neighbours(self) -> &'static [(DirectionalKeypad, Self)] {
160-
use DirectionalKeypad::*;
207+
const fn coords(self) -> Point2D<u32> {
208+
use NumericKeypad::*;
161209
match self {
162-
Up => &[(Right, Activate), (Down, Down)],
163-
Activate => &[(Left, Up), (Down, Right)],
164-
Left => &[(Right, Down)],
165-
Down => &[(Up, Up), (Left, Left), (Right, Right)],
166-
Right => &[(Up, Activate), (Left, Down)],
210+
Key7 => Point2D::new(0, 0),
211+
Key8 => Point2D::new(1, 0),
212+
Key9 => Point2D::new(2, 0),
213+
Key4 => Point2D::new(0, 1),
214+
Key5 => Point2D::new(1, 1),
215+
Key6 => Point2D::new(2, 1),
216+
Key1 => Point2D::new(0, 2),
217+
Key2 => Point2D::new(1, 2),
218+
Key3 => Point2D::new(2, 2),
219+
Key0 => Point2D::new(1, 3),
220+
Activate => Point2D::new(2, 3),
167221
}
168222
}
223+
224+
cost_matrix_functions!();
169225
}
170226

171227
examples!(Day21 -> (u64, u64) [

0 commit comments

Comments
 (0)