Rust/2024/06: improve solution

Rust/2024/06.rs

@@ -38,11 +38,10 @@ struct Walk<'a> {
     d: Direction,
     grid: &'a Input,
     init: bool,
-    extra_obstruction: Option<(usize, usize)>,
 }
 
 impl Walk<'_> {
-    fn from_input(input: &Input, extra_obstruction: Option<(usize, usize)>) -> Walk {
+    fn from_input(input: &Input) -> Walk {
         let (x, y) = input
             .iter()
             .enumerate()
@@ -54,19 +53,8 @@ impl Walk<'_> {
             d: Direction::North,
             grid: input,
             init: false,
-            extra_obstruction,
         }
     }
-
-    fn check_cycle(self) -> bool {
-        let mut visited = FxHashSet::default();
-        for p in self {
-            if !visited.insert(p) {
-                return true;
-            }
-        }
-        false
-    }
 }
 
 impl Iterator for Walk<'_> {
@@ -82,9 +70,7 @@ impl Iterator for Walk<'_> {
             let (nx, ny) = self
                 .d
                 .step(self.x, self.y, self.grid[0].len(), self.grid.len())?;
-            if !matches!(self.grid[ny][nx], Cell::Obstruction)
-                && self.extra_obstruction != Some((nx, ny))
-            {
+            if !matches!(self.grid[ny][nx], Cell::Obstruction) {
                 (self.x, self.y) = (nx, ny);
                 break;
             }
@@ -95,26 +81,142 @@ impl Iterator for Walk<'_> {
     }
 }
 
+#[derive(Debug)]
+struct Grid {
+    width: usize,
+    height: usize,
+    start: (usize, usize),
+    north: Vec<LookupRow>,
+    east: Vec<LookupRow>,
+    south: Vec<LookupRow>,
+    west: Vec<LookupRow>,
+}
+
+#[derive(Debug)]
+struct LookupRow(Vec<usize>);
+
+impl LookupRow {
+    fn lookup(&self, n: usize, extra: Option<usize>) -> usize {
+        match extra {
+            Some(x) if x <= n => self.0[n].min(n - x),
+            _ => self.0[n],
+        }
+    }
+}
+
+impl FromIterator<bool> for LookupRow {
+    fn from_iter<T: IntoIterator<Item = bool>>(iter: T) -> Self {
+        Self(
+            iter.into_iter()
+                .scan(usize::MAX, |s, x| {
+                    *s = if x { 0 } else { s.saturating_add(1) };
+                    Some(*s)
+                })
+                .collect(),
+        )
+    }
+}
+
+impl Grid {
+    fn from_input(input: &Input) -> Self {
+        let width = input[0].len();
+        let height = input.len();
+
+        let start = input
+            .iter()
+            .enumerate()
+            .find_map(|(y, row)| row.iter().index(&Cell::Guard).map(|x| (x, y)))
+            .unwrap();
+
+        let is_obstruction = |x: usize, y: usize| matches!(input[y][x], Cell::Obstruction);
+
+        Self {
+            width,
+            height,
+            start,
+            north: (0..width)
+                .into_par_iter()
+                .map(|x| (0..height).map(|y| is_obstruction(x, y)).collect())
+                .collect(),
+            east: (0..height)
+                .into_par_iter()
+                .map(|y| (0..width).rev().map(|x| is_obstruction(x, y)).collect())
+                .collect(),
+            south: (0..width)
+                .into_par_iter()
+                .map(|x| (0..height).rev().map(|y| is_obstruction(x, y)).collect())
+                .collect(),
+            west: (0..height)
+                .into_par_iter()
+                .map(|y| (0..width).map(|x| is_obstruction(x, y)).collect())
+                .collect(),
+        }
+    }
+
+    fn check_cycle(&self, extra_obstruction: (usize, usize)) -> bool {
+        let (mut x, mut y) = self.start;
+        let mut d = Direction::North;
+        let mut visited = FxHashSet::default();
+        while visited.insert((x, y, d)) {
+            let Some((nx, ny)) = self.walk(x, y, d, extra_obstruction) else {
+                return false;
+            };
+            (x, y, d) = (nx, ny, d.rotate_right())
+        }
+        true
+    }
+
+    fn walk(
+        &self,
+        x: usize,
+        y: usize,
+        direction: Direction,
+        extra_obstruction: (usize, usize),
+    ) -> Option<(usize, usize)> {
+        let eo = Some(extra_obstruction);
+        Some(match direction {
+            Direction::North => self.north[x].lookup(y, eo.filter(|eo| eo.0 == x).map(|eo| eo.1)),
+            Direction::East => self.east[y].lookup(
+                self.width - x - 1,
+                eo.filter(|eo| eo.1 == y).map(|eo| self.width - eo.0 - 1),
+            ),
+            Direction::South => self.south[x].lookup(
+                self.height - y - 1,
+                eo.filter(|eo| eo.0 == x).map(|eo| self.height - eo.1 - 1),
+            ),
+            Direction::West => self.west[y].lookup(x, eo.filter(|eo| eo.1 == y).map(|eo| eo.0)),
+        })
+        .filter(|&r| r != usize::MAX)
+        .map(|r| {
+            let r = r - 1;
+            match direction {
+                Direction::North => (x, y - r),
+                Direction::East => (x + r, y),
+                Direction::South => (x, y + r),
+                Direction::West => (x - r, y),
+            }
+        })
+    }
+}
+
 fn part1(input: &Input) -> usize {
-    Walk::from_input(input, None)
+    Walk::from_input(input)
         .map(|(x, y, _)| (x, y))
         .collect::<FxHashSet<_>>()
         .len()
 }
 
 fn part2(input: &Input) -> usize {
-    let candidates = Walk::from_input(input, None)
-        .flat_map(|(x, y, d)| {
-            [d, d.rotate_right()]
-                .into_iter()
-                .flat_map(move |d| d.step(x, y, input[0].len(), input.len()))
-        })
+    let candidates = Walk::from_input(input)
+        .map(|(x, y, _)| (x, y))
         .filter(|&(x, y)| matches!(input[y][x], Cell::Empty))
         .collect::<FxHashSet<_>>();
 
+    let grid = Grid::from_input(input);
+
     candidates
         .into_par_iter()
-        .filter(|&(ox, oy)| Walk::from_input(input, Some((ox, oy))).check_cycle())
+        .filter(|&(ox, oy)| grid.check_cycle((ox, oy)))
         .count()
 }