1use utils::geometry::Vec2;
2use utils::prelude::*;
3
4#[derive(Clone, Debug)]
6pub struct Day21 {
7    codes: Vec<u16>,
8}
9
10utils::enumerable_enum! {
11    #[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord)]
12    enum NumericKeypad {
13        Key7 = 7, Key8 = 8, Key9 = 9,
14        Key4 = 4, Key5 = 5, Key6 = 6,
15        Key1 = 1, Key2 = 2, Key3 = 3,
16                  Key0 = 0, Activate = 10,
17    }
18}
19
20utils::enumerable_enum! {
21    #[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord)]
22    enum DirectionalKeypad {
23                  Up   = 0, Activate = 4,
24        Left = 1, Down = 2, Right = 3,
25    }
26}
27
28#[cfg(feature = "const_lut")]
29static PART1_MATRIX: [[u64; 11]; 11] = num_matrix(2);
30#[cfg(feature = "const_lut")]
31static PART2_MATRIX: [[u64; 11]; 11] = num_matrix(25);
32
33impl Day21 {
34    pub fn new(input: &str, _: InputType) -> Result<Self, InputError> {
35        Ok(Self {
36            codes: parser::number_range(0..=999)
37                .with_suffix(b'A')
38                .parse_lines(input)?,
39        })
40    }
41
42    #[must_use]
43    pub fn part1(&self) -> u64 {
44        #[cfg(feature = "const_lut")]
45        return self.complexity(&PART1_MATRIX);
46
47        #[cfg(not(feature = "const_lut"))]
48        return self.complexity(&num_matrix(2));
49    }
50
51    #[must_use]
52    pub fn part2(&self) -> u64 {
53        #[cfg(feature = "const_lut")]
54        return self.complexity(&PART2_MATRIX);
55
56        #[cfg(not(feature = "const_lut"))]
57        return self.complexity(&num_matrix(25));
58    }
59
60    fn complexity(&self, matrix: &[[u64; 11]; 11]) -> u64 {
61        self.codes
62            .iter()
63            .map(|&code| {
64                let digits = [(code % 1000) / 100, (code % 100) / 10, code % 10];
65                let length = matrix[NumericKeypad::Activate as usize][digits[0] as usize]
66                    + matrix[digits[0] as usize][digits[1] as usize]
67                    + matrix[digits[1] as usize][digits[2] as usize]
68                    + matrix[digits[2] as usize][NumericKeypad::Activate as usize];
69                length * code as u64
70            })
71            .sum()
72    }
73}
74
75const fn num_matrix(robots: u32) -> [[u64; NumericKeypad::COUNT]; NumericKeypad::COUNT] {
76    let mut dir_matrix = [[1; 5]; 5];
77    let mut i = 0;
78    while i < robots {
79        dir_matrix = DirectionalKeypad::cost_matrix(dir_matrix);
80        i += 1;
81    }
82    NumericKeypad::cost_matrix(dir_matrix)
83}
84
85macro_rules! cost_matrix_functions {
87    () => {
88        const fn cost_matrix(
89            dir_matrix: [[u64; DirectionalKeypad::COUNT]; DirectionalKeypad::COUNT],
90        ) -> [[u64; Self::COUNT]; Self::COUNT] {
91            let mut result = [[u64::MAX; Self::COUNT]; Self::COUNT];
92            let mut i = 0;
93            while i < Self::COUNT {
94                result[i][i] = 1;
95                Self::visit(
96                    0,
97                    Self::ALL[i],
98                    DirectionalKeypad::Activate,
99                    Self::ALL[i],
100                    dir_matrix,
101                    &mut result,
102                );
103                i += 1;
104            }
105            result
106        }
107
108        const fn visit(
109            cost: u64,
110            current: Self,
111            parent: DirectionalKeypad,
112            start: Self,
113            dir_matrix: [[u64; DirectionalKeypad::COUNT]; DirectionalKeypad::COUNT],
114            result: &mut [[u64; Self::COUNT]; Self::COUNT],
115        ) {
116            let cost_with_activate =
117                cost + dir_matrix[parent as usize][DirectionalKeypad::Activate as usize];
118            if cost_with_activate < result[start as usize][current as usize] {
119                result[start as usize][current as usize] = cost_with_activate;
120            }
121
122            let start_coords = start.coords();
123            let current_coords = current.coords();
124            let current_distance = current_coords.x.abs_diff(start_coords.x)
125                + current_coords.y.abs_diff(start_coords.y);
126
127            let neighbours = current.neighbours();
128            let mut i = 0;
129            while i < neighbours.len() {
130                let (next_parent, next) = neighbours[i];
131                let next_point = next.coords();
132                let next_distance =
133                    next_point.x.abs_diff(start_coords.x) + next_point.y.abs_diff(start_coords.y);
134                if next_distance > current_distance {
135                    Self::visit(
136                        cost + dir_matrix[parent as usize][next_parent as usize],
137                        next,
138                        next_parent,
139                        start,
140                        dir_matrix,
141                        result,
142                    );
143                }
144                i += 1;
145            }
146        }
147    };
148}
149
150impl DirectionalKeypad {
151    const fn neighbours(self) -> &'static [(DirectionalKeypad, Self)] {
152        use DirectionalKeypad::*;
153        match self {
154            Up => &[(Right, Activate), (Down, Down)],
155            Activate => &[(Left, Up), (Down, Right)],
156            Left => &[(Right, Down)],
157            Down => &[(Up, Up), (Left, Left), (Right, Right)],
158            Right => &[(Up, Activate), (Left, Down)],
159        }
160    }
161
162    const fn coords(self) -> Vec2<u32> {
163        use DirectionalKeypad::*;
164        match self {
165            Up => Vec2::new(1, 0),
166            Activate => Vec2::new(2, 0),
167            Left => Vec2::new(0, 1),
168            Down => Vec2::new(1, 1),
169            Right => Vec2::new(2, 1),
170        }
171    }
172
173    cost_matrix_functions!();
174}
175
176impl NumericKeypad {
177    const fn neighbours(self) -> &'static [(DirectionalKeypad, Self)] {
178        use DirectionalKeypad::{Down, Left, Right, Up};
179        use NumericKeypad::*;
180        match self {
181            Key7 => &[(Right, Key8), (Down, Key4)],
182            Key8 => &[(Left, Key7), (Right, Key9), (Down, Key5)],
183            Key9 => &[(Left, Key8), (Down, Key6)],
184            Key4 => &[(Up, Key7), (Right, Key5), (Down, Key1)],
185            Key5 => &[(Up, Key8), (Left, Key4), (Right, Key6), (Down, Key2)],
186            Key6 => &[(Up, Key9), (Left, Key5), (Down, Key3)],
187            Key1 => &[(Up, Key4), (Right, Key2)],
188            Key2 => &[(Up, Key5), (Left, Key1), (Right, Key3), (Down, Key0)],
189            Key3 => &[(Up, Key6), (Left, Key2), (Down, Activate)],
190            Key0 => &[(Up, Key2), (Right, Activate)],
191            Activate => &[(Up, Key3), (Left, Key0)],
192        }
193    }
194
195    const fn coords(self) -> Vec2<u32> {
196        use NumericKeypad::*;
197        match self {
198            Key7 => Vec2::new(0, 0),
199            Key8 => Vec2::new(1, 0),
200            Key9 => Vec2::new(2, 0),
201            Key4 => Vec2::new(0, 1),
202            Key5 => Vec2::new(1, 1),
203            Key6 => Vec2::new(2, 1),
204            Key1 => Vec2::new(0, 2),
205            Key2 => Vec2::new(1, 2),
206            Key3 => Vec2::new(2, 2),
207            Key0 => Vec2::new(1, 3),
208            Activate => Vec2::new(2, 3),
209        }
210    }
211
212    cost_matrix_functions!();
213}
214
215examples!(Day21 -> (u64, u64) [
216    {input: "029A\n980A\n179A\n456A\n379A", part1: 126384},
217]);