Optimize 2024 day 20

Rewriting to allow vectorization reduces runtime from ~15ms to ~4ms, or ~1ms with target_cpu=native.

Author: ictrobot

crates/year2024/src/day20.rs

@@ -5,11 +5,13 @@ use utils::prelude::*;
 /// Finding shortcuts phasing through walls in a maze.
 #[derive(Clone, Debug)]
 pub struct Day20 {
-    grid: Vec<u8>,
-    distances: Vec<u32>,
+    distances: Vec<Distance>,
     cols: usize,
 }
 
+// Depending on AVX support, u16 or u32 can be faster
+type Distance = u16;
+
 impl Day20 {
     pub fn new(input: &str, _: InputType) -> Result<Self, InputError> {
         let (_, cols, mut grid) = grid::from_str_padded(input, 20, b'#', |b| match b {
@@ -35,58 +37,46 @@ impl Day20 {
         }
         grid[end] = b'.';
 
-        let mut distances = vec![u32::MAX; grid.len()];
+        let mut distances = vec![Distance::MAX; grid.len()];
         distances[start] = 0;
         let mut queue = VecDeque::new();
-        queue.push_back((start, 0));
+        queue.push_back((start, 0 as Distance));
         while let Some((index, distance)) = queue.pop_front() {
             if index == end {
                 break;
             }
 
+            let Some(next_distance) = distance.checked_add(1) else {
+                return Err(InputError::new(input, 0, "path too long"));
+            };
+
             for offset in [1, cols as isize, -1, -(cols as isize)] {
                 let next = index.wrapping_add_signed(offset);
-                if grid[next] == b'.' && distances[next] == u32::MAX {
-                    distances[next] = distance + 1;
-                    queue.push_back((next, distance + 1));
+                if grid[next] == b'.' && distances[next] == Distance::MAX {
+                    distances[next] = next_distance;
+                    queue.push_back((next, next_distance));
                 }
             }
         }
 
-        Ok(Self {
-            grid,
-            distances,
-            cols,
-        })
+        Ok(Self { distances, cols })
     }
 
     #[must_use]
     pub fn part1(&self) -> u32 {
         let mut cheats = 0;
-        for index in (self.cols * 20)..self.grid.len() - (self.cols * 20) {
-            let this_distance = self.distances[index];
-            if this_distance == u32::MAX {
-                continue;
-            }
-
-            for offset in [1, self.cols as isize, -1, -(self.cols as isize)] {
-                if self.grid[index.wrapping_add_signed(offset)] != b'#' {
-                    continue;
-                }
-
-                let Some(next) = index.checked_add_signed(offset * 2) else {
-                    continue;
-                };
-                let Some(&target_distance) = self.distances.get(next) else {
-                    continue;
-                };
-                if target_distance == u32::MAX || target_distance < this_distance + 2 {
-                    continue;
-                }
-                let diff = target_distance - this_distance - 2;
-                if diff >= 100 {
-                    cheats += 1;
-                }
+        for offset in [2, self.cols as isize * 2, -2, self.cols as isize * -2] {
+            for (index, target) in (self.cols * 20..self.distances.len() - (self.cols * 20))
+                .zip((self.cols * 20).wrapping_add_signed(offset)..)
+            {
+                let this_distance = self.distances[index];
+                let target_distance = self.distances[target];
+                cheats += u32::from(
+                    (target_distance != Distance::MAX)
+                        & (this_distance != Distance::MAX)
+                        & (target_distance > this_distance.wrapping_add(2))
+                        & (target_distance.wrapping_sub(this_distance).wrapping_sub(2) >= 100),
+                );
             }
         }
         cheats
@@ -95,36 +85,32 @@ impl Day20 {
     #[must_use]
     pub fn part2(&self) -> u32 {
         let mut cheats = 0;
-        for index in (self.cols * 20)..self.grid.len() - (self.cols * 20) {
-            let this_distance = self.distances[index];
-            if this_distance == u32::MAX {
-                continue;
-            }
-
-            for x_offset in -20isize..=20 {
-                let y_limit = 20 - x_offset.abs();
-                for y_offset in -y_limit..=y_limit {
-                    let cheat_length = (x_offset.unsigned_abs() + y_offset.unsigned_abs()) as u32;
-                    if cheat_length == 0 {
-                        continue;
-                    }
-
-                    let offset = y_offset * (self.cols as isize) + x_offset;
-                    let next = index.wrapping_add_signed(offset);
-                    let target_distance = self.distances[next];
-                    if target_distance == u32::MAX || target_distance < this_distance + cheat_length
-                    {
-                        continue;
-                    }
+        for x_offset in -20isize..=20 {
+            let y_limit = 20 - x_offset.abs();
+            for y_offset in -y_limit..=y_limit {
+                let cheat_length = (x_offset.unsigned_abs() + y_offset.unsigned_abs()) as Distance;
+                if cheat_length == 0 {
+                    continue;
+                }
 
-                    let diff = target_distance - this_distance - cheat_length;
-                    if diff >= 100 {
-                        cheats += 1;
-                    }
+                let offset = y_offset * (self.cols as isize) + x_offset;
+                for (index, target) in (self.cols * 20..self.distances.len() - (self.cols * 20))
+                    .zip((self.cols * 20).wrapping_add_signed(offset)..)
+                {
+                    let this_distance = self.distances[index];
+                    let target_distance = self.distances[target];
+                    cheats += u32::from(
+                        (target_distance != Distance::MAX)
+                            & (this_distance != Distance::MAX)
+                            & (target_distance > this_distance.wrapping_add(cheat_length))
+                            & (target_distance
+                                .wrapping_sub(this_distance)
+                                .wrapping_sub(cheat_length)
+                                >= 100),
+                    );
                 }
             }
         }
-
         cheats
     }
 }