Simplify y-hit iterator

2021/src/day17.rs

@@ -1,8 +1,6 @@
 use std::io::Read;
 use std::ops::RangeInclusive;
 
-use itertools::Itertools;
-use itertools::MinMaxResult;
 use nom::bytes::complete::tag;
 use nom::combinator::map;
 use nom::sequence::preceded;
@@ -19,78 +17,61 @@ fn solve_quadratic(a: f64, b: f64, c: f64) -> Option<f64> {
         None
     } else {
         // Don't care about the smaller solution due to problem statement
+        if a > 0. {
+            Some((-b + d.sqrt()) / 2. / a)
+        } else {
             Some((-b - d.sqrt()) / 2. / a)
         }
     }
+}
 
 fn position(initial: i32, time: i32) -> i32 {
     time * (2 * initial - time + 1) / 2
 }
 
-fn find_hit(initial: i32, range: &RangeInclusive<i32>) -> Option<RangeInclusive<i32>> {
+fn find_hit(initial: i32, range: &RangeInclusive<i32>) -> impl Iterator<Item = i32> + '_ {
     // y position at time x: f(x) = x * (1 + initial + initial - x) / 2
     //                            = -1/2x^2 + (initial + 0.5)x
     //
     // to hit, find x := (max(box) + min(box)) / 2 = f(x)
     //                                             = -1/2x^2 + (initial + 0.5)x
     // -1/2x^2 + (initial + 0.5)x - (max(box) + min(box)) / 2 = 0
-    let middle = (*range.start() + *range.end()) as f64 / 2.;
+    let middle = *range.start() as f64;
     let b = initial as f64 + 0.5;
-    let hit = solve_quadratic(-0.5, b, -middle)? as i32;
-
-    if hit < 0 {
-        // Should not happen because of the shape but for correctness
-        None
+    let hit = if let Some(hit) = solve_quadratic(-0.5, b, -middle) {
+        hit as i32
     } else {
-        let min_hit = (0..=hit)
+        // Cause an empty range
+        -1
+    };
+
+    (0..=hit)
         .rev()
-            .take_while(|&n| range.contains(&position(initial, n)))
-            .min();
-
-        let max_hit = ((hit + 1)..)
-            .take_while(|&n| range.contains(&position(initial, n)))
-            .max();
-
-        match (min_hit, max_hit) {
-            (Some(min), Some(max)) => Some(min..=max),
-            (Some(val), None) | (None, Some(val)) => Some(val..=val),
-            _ => None,
-        }
-    }
+        .take_while(move |&n| range.contains(&position(initial, n)))
 }
 
-fn find_speed(x: i32, range: &RangeInclusive<i32>) -> Option<RangeInclusive<i32>> {
-    let mut min = 0;
-    let mut max = *range.end();
+fn find_speed(x: i32, range: &RangeInclusive<i32>) -> Option<(i32, i32)> {
+    if *range.end() <= position(x, x) {
+        // Can and should come to a full stop
+        let max = solve_quadratic(0.5, 0.5, -*range.end() as f64)? as i32;
 
-    // Need to tweak the formula as x slows down
-    let x_pos = |speed| position(speed, speed.min(x));
-
-    while max >= min {
-        let speed = (max + min) / 2;
-
-        let pos = x_pos(speed);
-
-        if range.contains(&x_pos(speed)) {
-            let min_speed = (0..speed)
+        let min = (0..=max)
             .rev()
-                .take_while(|&speed| range.contains(&x_pos(speed)))
-                .min()
-                .unwrap_or(speed);
+            .take_while(|&n| range.contains(&position(n, n)))
+            .last()?;
 
-            let max_speed = ((speed + 1)..)
-                .take_while(|&speed| range.contains(&x_pos(speed)))
-                .max()
-                .unwrap_or(speed);
-            return Some(min_speed..=max_speed);
-        } else if pos < *range.start() {
-            min = speed + 1;
+        Some((min, max))
     } else {
-            max = speed - 1;
-        }
-    }
+        // Might hit the target at speed
+        let max = (x * x + 2 * *range.end() - x) / (2 * x);
 
-    None
+        let min = (0..=max)
+            .rev()
+            .take_while(|&n| range.contains(&position(n, n.min(x))))
+            .last()?;
+
+        Some((min, max))
+    }
 }
 
 fn parse_range(input: &[u8]) -> IResult<&[u8], RangeInclusive<i32>> {
@@ -111,40 +92,35 @@ fn parse_input(input: &[u8]) -> IResult<&[u8], (RangeInclusive<i32>, RangeInclus
 pub fn part1(input: &mut dyn Read) -> String {
     let (x_range, y_range) = read_input(input, parse_input);
 
-    let check_value = |y_speed| {
-        let mut time = find_hit(y_speed, &y_range)?;
+    let check_value =
+        |y_speed| find_hit(y_speed, &y_range).any(|time| find_speed(time, &x_range).is_some());
 
-        if time.any(|time| find_speed(time, &x_range).is_some()) {
-            Some(position(y_speed, y_speed))
-        } else {
-            None
-        }
-    };
+    debug_assert!(*y_range.start() < 0);
+    let y_max = -*y_range.start();
 
-    (0..1000).filter_map(check_value).max().unwrap().to_string()
+    (0..y_max)
+        .rev()
+        .find(|&speed| check_value(speed))
+        .map(|speed| position(speed, speed))
+        .unwrap()
+        .to_string()
 }
 
 pub fn part2(input: &mut dyn Read) -> String {
     let (x_range, y_range) = read_input(input, parse_input);
 
     let num_options = |y_speed| {
-        let time = find_hit(y_speed, &y_range)?;
-        let range = time
+        find_hit(y_speed, &y_range)
             .filter_map(|time| find_speed(time, &x_range))
-            .flat_map(|x| [*x.start(), *x.end()])
-            .minmax();
-
-        Some(match range {
-            MinMaxResult::NoElements => 0,
-            MinMaxResult::OneElement(_) => 1,
-            MinMaxResult::MinMax(min, max) => max - min + 1,
-        })
+            .reduce(|(a_min, a_max), (b_min, b_max)| (a_min.min(b_min), a_max.max(b_max)))
+            .map(|(min, max)| max - min + 1)
+            .unwrap_or(0)
     };
 
-    (-1000..1000)
-        .filter_map(num_options)
-        .sum::<i32>()
-        .to_string()
+    debug_assert!(*y_range.start() < 0);
+    let y_max = -*y_range.start();
+
+    (-y_max..y_max).map(num_options).sum::<i32>().to_string()
 }
 
 #[cfg(test)]
@@ -157,9 +133,9 @@ mod tests {
 
     #[test]
     fn test_find_hit() {
-        assert_eq!(find_hit(2, &(-10..=-5)), Some(7..=7));
-        assert_eq!(find_hit(3, &(-10..=-5)), Some(9..=9));
-        assert_eq!(find_hit(0, &(-10..=-5)), Some(4..=5));
+        assert_eq!(find_hit(2, &(-10..=-5)).collect::<Vec<_>>(), vec![7]);
+        assert_eq!(find_hit(3, &(-10..=-5)).collect::<Vec<_>>(), vec![9]);
+        assert_eq!(find_hit(0, &(-10..=-5)).collect::<Vec<_>>(), vec![5, 4]);
     }
 
     #[test]