Slight simplification

ephemient · Dec 23, 2023 · df1f58d · df1f58d
1 parent 8a534e3
commit df1f58d
Showing 1 changed file with 33 additions and 49 deletions.
diff --git a/rs/src/day21.rs b/rs/src/day21.rs
@@ -1,4 +1,5 @@
 use std::collections::BTreeSet;
+use std::iter;
 
 #[allow(clippy::type_complexity)]
 fn parse(data: &str) -> Option<(Vec<&[u8]>, (usize, usize))> {
@@ -12,63 +13,46 @@ fn parse(data: &str) -> Option<(Vec<&[u8]>, (usize, usize))> {
 }
 
 fn dec((m, r): (isize, usize), n: usize) -> (isize, usize) {
-    if let Some(r) = r.checked_sub(1) {
-        (m, r)
-    } else {
-        (m - 1, n - 1)
-    }
+    r.checked_sub(1).map_or((m - 1, n - 1), |r| (m, r))
 }
 
 fn inc((m, r): (isize, usize), n: usize) -> (isize, usize) {
-    if let Some(r) = r.checked_add(1).filter(|r| *r < n) {
-        (m, r)
-    } else {
-        (m + 1, 0)
-    }
+    r.checked_add(1)
+        .filter(|r| *r < n)
+        .map_or((m + 1, 0), |r| (m, r))
 }
 
-fn bfs<'a>(
-    grid: &'a [&'a [u8]],
-    start: (usize, usize),
-    n: impl IntoIterator<Item = usize> + 'a,
-) -> impl Iterator<Item = usize> + 'a {
-    n.into_iter().scan(
-        (
-            0,
-            0,
-            0,
-            BTreeSet::<((isize, usize), (isize, usize))>::new(),
-            BTreeSet::from([((0isize, start.0), (0isize, start.1))]),
-        ),
-        |(even, odd, d, v, q), n| {
-            while *d <= n {
-                *if *d & 1 == 0 { &mut *even } else { &mut *odd } += q.len();
-                *d += 1;
-                v.extend(q.iter());
-                *q = q
-                    .iter()
-                    .flat_map(|p| {
-                        let h = grid.len();
-                        let w = grid[p.0 .1].len();
-                        [
-                            (dec(p.0, h), p.1),
-                            (p.0, dec(p.1, w)),
-                            (p.0, inc(p.1, w)),
-                            (inc(p.0, h), p.1),
-                        ]
-                    })
-                    .filter(|p| !v.contains(p) && grid[p.0 .1][p.1 .1] != b'#')
-                    .collect();
-            }
-            Some(*if n & 1 == 0 { even } else { odd })
-        },
-    )
+fn bfs<'a>(grid: &'a [&'a [u8]], start: (usize, usize)) -> impl Iterator<Item = usize> + 'a {
+    let (mut d, mut even, mut odd) = (0, 0, 0);
+    let mut v = BTreeSet::<((isize, usize), (isize, usize))>::new();
+    let mut q = BTreeSet::from([((0, start.0), (0, start.1))]);
+    iter::from_fn(move || {
+        let n = if d & 1 == 0 { &mut even } else { &mut odd };
+        *n += q.len();
+        d += 1;
+        v.extend(q.iter());
+        q = q
+            .iter()
+            .flat_map(|p| {
+                let h = grid.len();
+                let w = grid[p.0 .1].len();
+                [
+                    (dec(p.0, h), p.1),
+                    (p.0, dec(p.1, w)),
+                    (p.0, inc(p.1, w)),
+                    (inc(p.0, h), p.1),
+                ]
+            })
+            .filter(|p| !v.contains(p) && grid[p.0 .1][p.1 .1] != b'#')
+            .collect();
+        Some(*n)
+    })
 }
 
 fn part1_n(data: &str, n: usize) -> Option<usize> {
     let (grid, start) = parse(data)?;
-    let mut iter = bfs(&grid, start, [n]);
-    iter.next()
+    let mut iter = bfs(&grid, start);
+    iter.nth(n)
 }
 
 pub fn part1(data: &str) -> Option<usize> {
@@ -79,7 +63,7 @@ fn part2_n(data: &str, n: usize) -> Option<usize> {
     let (grid, start) = parse(data)?;
     let m = grid.len();
     let (q, r) = (n / m, n % m);
-    let mut iter = bfs(&grid, start, [r, r + m, r + 2 * m, r + 3 * m]);
+    let mut iter = bfs(&grid, start).skip(r).step_by(m);
     let (a, b, c, d) = (iter.next()?, iter.next()?, iter.next()?, iter.next()?);
     if a + 3 * (c - b) == d {
         Some(a + (b - a) * q + (a + c - 2) * (q * q.saturating_sub(1) / 2))