fix: replace hand-rolled quickselect with select_nth_unstable_by

benches/order_statistics.rs

@@ -94,5 +94,66 @@ fn bench_order_statistic(c: &mut Criterion) {
     group.finish();
 }
 
-criterion_group!(benches, bench_order_statistic);
+/// Selection cost across input sizes and shapes.
+///
+/// Kept separate from the fixed-size group above so the two can evolve
+/// independently, and written against only the public `Data` API so the same
+/// file compiles on any branch. To compare two implementations:
+///
+/// ```text
+/// git checkout main    && cargo bench --bench order_statistics -- --save-baseline main
+/// git checkout <branch> && cargo bench --bench order_statistics -- --baseline main
+/// ```
+///
+/// The shapes matter as much as the sizes. A median-of-three pivot is fine on
+/// random data and degrades on input built to defeat it, so a random-only
+/// benchmark would understate the difference between a quickselect and an
+/// introselect with an O(n) fallback.
+fn bench_selection_scaling(c: &mut Criterion) {
+    fn shaped(shape: &str, n: usize) -> Vec<f64> {
+        let mut rng = StdRng::seed_from_u64(0xB0A7);
+        match shape {
+            // shuffled, the ordinary case
+            "random" => {
+                let mut v: Vec<f64> = (0..n).map(|i| i as f64).collect();
+                v.shuffle(&mut rng);
+                v
+            }
+            // already ordered, both directions
+            "sorted" => (0..n).map(|i| i as f64).collect(),
+            "reversed" => (0..n).rev().map(|i| i as f64).collect(),
+            // ascends then descends: the classic median-of-three adversary,
+            // since the first, middle and last elements are unrepresentative
+            "organ_pipe" => (0..n)
+                .map(|i| if i < n / 2 { i } else { n - i } as f64)
+                .collect(),
+            // few distinct values, so partitions are heavily unbalanced
+            "duplicates" => {
+                let distinct = ((n as f64).sqrt() as usize).max(1);
+                let mut v: Vec<f64> = (0..n).map(|i| (i % distinct) as f64).collect();
+                v.shuffle(&mut rng);
+                v
+            }
+            other => panic!("unknown shape {other}"),
+        }
+    }
+
+    let mut group = c.benchmark_group("selection scaling");
+    for &n in &[1_024usize, 65_536, 1_048_576] {
+        for shape in ["random", "sorted", "reversed", "organ_pipe", "duplicates"] {
+            let data = shaped(shape, n);
+            group.throughput(Throughput::Elements(n as u64));
+            group.bench_function(format!("median/{shape}/{n}"), |b| {
+                b.iter_batched(
+                    || Data::new(data.clone()),
+                    |data| black_box(data.median()),
+                    BatchSize::LargeInput,
+                )
+            });
+        }
+    }
+    group.finish();
+}
+
+criterion_group!(benches, bench_order_statistic, bench_selection_scaling);
 criterion_main!(benches);

src/statistics/slice_statistics.rs

@@ -64,8 +64,19 @@ impl<D: AsMut<[f64]> + AsRef<[f64]>> Data<D> {
         self.0.as_ref().iter()
     }
 
-    // Selection algorithm from Numerical Recipes
-    // See: https://en.wikipedia.org/wiki/Selection_algorithm
+    // Selection via `<[T]>::select_nth_unstable_by`, an introselect with an
+    // O(n) worst case. `total_cmp` also gives NaN a defined position (positive
+    // NaNs sort past +inf) rather than the arbitrary result the previous
+    // Numerical Recipes partition produced, which is the point of the change.
+    //
+    // The ordered fast paths are not incidental. The NR quickselect took its
+    // pivot from a median of three, which on sorted or reverse-sorted input
+    // lands on the true median immediately and finishes in a single partition
+    // -- so it beat plain introselect on exactly those shapes. Testing for
+    // order first recovers that, and improves on it: the answer becomes an
+    // index rather than a partition. Both checks stop at the first out-of-order
+    // pair, costing a couple of comparisons on unordered input, and they also
+    // pay off in `median` and `quantile`, which call this twice.
     fn select_inplace(&mut self, rank: usize) -> f64 {
         if rank == 0 {
             return self.min();
@@ -74,61 +85,19 @@ impl<D: AsMut<[f64]> + AsRef<[f64]>> Data<D> {
             return self.max();
         }
 
-        let mut low = 0;
-        let mut high = self.len() - 1;
-        loop {
-            if high <= low + 1 {
-                if high == low + 1 && self[high] < self[low] {
-                    self.swap(low, high)
-                }
-                return self[rank];
-            }
-
-            let middle = (low + high) / 2;
-            self.swap(middle, low + 1);
-
-            if self[low] > self[high] {
-                self.swap(low, high);
-            }
-            if self[low + 1] > self[high] {
-                self.swap(low + 1, high);
-            }
-            if self[low] > self[low + 1] {
-                self.swap(low, low + 1);
-            }
-
-            let mut begin = low + 1;
-            let mut end = high;
-            let pivot = self[begin];
-            loop {
-                loop {
-                    begin += 1;
-                    if self[begin] >= pivot {
-                        break;
-                    }
-                }
-                loop {
-                    end -= 1;
-                    if self[end] <= pivot {
-                        break;
-                    }
-                }
-                if end < begin {
-                    break;
-                }
-                self.swap(begin, end);
-            }
-
-            self[low + 1] = self[end];
-            self[end] = pivot;
-
-            if end >= rank {
-                high = end - 1;
-            }
-            if end <= rank {
-                low = begin;
-            }
+        // `<=` rather than `total_cmp` here, which is both cheaper per element
+        // and still correct: any NaN makes some comparison false, so a slice
+        // containing one never takes a fast path and falls through to the
+        // general case, and without NaN the two orderings agree.
+        let slice = self.0.as_mut();
+        if slice.is_sorted_by(|a, b| a <= b) {
+            return slice[rank];
+        }
+        if slice.is_sorted_by(|a, b| a >= b) {
+            return slice[slice.len() - 1 - rank];
         }
+
+        *slice.select_nth_unstable_by(rank, |a, b| a.total_cmp(b)).1
     }
 }
 
@@ -550,6 +519,141 @@ mod tests {
 
     // TODO: test codeplex issue 5667 (Math.NET)
 
+    /// Selection agrees with sorting, on ordinary values across a range of
+    /// shapes and sizes.
+    ///
+    /// The counterpart to the NaN and infinity tests below, which pin that
+    /// nothing panics but say nothing about the values returned. This also
+    /// covers the ordered fast paths in `select_inplace`: those return an
+    /// index instead of partitioning, so they have to produce exactly what the
+    /// general path would, including for the descending case where the index
+    /// is counted from the far end.
+    #[test]
+    #[cfg(feature = "std")]
+    fn test_order_statistics_match_sorted_reference() {
+        fn shaped(shape: &str, n: usize) -> Vec<f64> {
+            match shape {
+                // deterministic pseudo-random, no rand dependency needed
+                "scattered" => (0..n).map(|i| ((i * 7919 + 13) % 1000) as f64 * 0.5).collect(),
+                "sorted" => (0..n).map(|i| i as f64 * 1.5).collect(),
+                "reversed" => (0..n).rev().map(|i| i as f64 * 1.5).collect(),
+                "organ_pipe" => (0..n).map(|i| if i < n / 2 { i } else { n - i } as f64).collect(),
+                "duplicates" => (0..n).map(|i| (i % 3) as f64).collect(),
+                "constant" => vec![4.25; n],
+                "negatives" => (0..n).map(|i| -((i * 37 % 100) as f64)).collect(),
+                other => panic!("unknown shape {other}"),
+            }
+        }
+
+        for n in [1usize, 2, 3, 4, 5, 8, 17, 101, 1000] {
+            for shape in [
+                "scattered", "sorted", "reversed", "organ_pipe", "duplicates", "constant",
+                "negatives",
+            ] {
+                let data = shaped(shape, n);
+                let mut sorted = data.clone();
+                sorted.sort_by(f64::total_cmp);
+
+                // every order statistic, 1-based
+                for order in 1..=n {
+                    let mut d = Data::new(data.clone());
+                    assert_eq!(
+                        d.order_statistic(order),
+                        sorted[order - 1],
+                        "{shape}/{n}: order_statistic({order})"
+                    );
+                }
+
+                // out of range on either side
+                let mut d = Data::new(data.clone());
+                assert!(d.order_statistic(0).is_nan(), "{shape}/{n}: order 0");
+                assert!(d.order_statistic(n + 1).is_nan(), "{shape}/{n}: order n+1");
+
+                // median, including the even case that averages two selections
+                let expected_median = if n % 2 == 1 {
+                    sorted[n / 2]
+                } else {
+                    (sorted[n / 2 - 1] + sorted[n / 2]) / 2.0
+                };
+                let mut d = Data::new(data.clone());
+                assert_eq!(
+                    OrderStatistics::median(&mut d),
+                    expected_median,
+                    "{shape}/{n}: median"
+                );
+
+                // the quantile endpoints, and that interior values stay in range
+                let mut d = Data::new(data.clone());
+                assert_eq!(d.quantile(0.0), sorted[0], "{shape}/{n}: quantile(0)");
+                assert_eq!(d.quantile(1.0), sorted[n - 1], "{shape}/{n}: quantile(1)");
+                for &tau in &[0.1, 0.25, 0.5, 0.75, 0.9] {
+                    let q = d.quantile(tau);
+                    assert!(
+                        q >= sorted[0] && q <= sorted[n - 1],
+                        "{shape}/{n}: quantile({tau}) = {q} outside [{}, {}]",
+                        sorted[0],
+                        sorted[n - 1]
+                    );
+                }
+            }
+        }
+    }
+
+    /// Selection must not depend on the order the input happens to arrive in.
+    #[test]
+    #[cfg(feature = "std")]
+    fn test_order_statistic_is_permutation_invariant() {
+        let base: Vec<f64> = (0..64).map(|i| ((i * 31 + 7) % 41) as f64).collect();
+        let mut sorted = base.clone();
+        sorted.sort_by(f64::total_cmp);
+
+        // a few fixed permutations, including the ones that hit the fast paths
+        let mut rotated = base.clone();
+        rotated.rotate_left(23);
+        let mut ascending = base.clone();
+        ascending.sort_by(f64::total_cmp);
+        let mut descending = ascending.clone();
+        descending.reverse();
+
+        for variant in [base.clone(), rotated, ascending, descending] {
+            for order in 1..=variant.len() {
+                let mut d = Data::new(variant.clone());
+                assert_eq!(d.order_statistic(order), sorted[order - 1]);
+            }
+        }
+    }
+
+    /// `select_inplace` used to implement the Numerical Recipes quickselect by
+    /// hand, whose partition loop walks off the end of the array when the slice
+    /// contains NaN (statrs-dev/statrs#163). `total_cmp` gives NaN a defined
+    /// place in a total order, so the search stays in bounds.
+    #[test]
+    fn test_order_statistics_with_nan_do_not_panic() {
+        // Arrays rather than `Vec`, so the test also builds under `no_std`.
+        let mut v: [f64; 1000] = core::array::from_fn(|i| (i as f64 * 0.7).sin());
+        for i in (0..1000).step_by(7) {
+            v[i] = f64::NAN;
+        }
+        let mut d = Data::new(v);
+        // the values themselves are unspecified once NaN is present; what is
+        // being pinned is that these terminate in bounds
+        let _ = d.quantile(0.5);
+        let _ = d.quantile(0.0);
+        let _ = d.quantile(1.0);
+        let _ = d.order_statistic(500);
+        let _ = OrderStatistics::median(&mut d);
+        let _ = d.interquartile_range();
+
+        // all-NaN, and NaN at the exact positions the old partition tripped on
+        let mut all_nan = Data::new([f64::NAN; 64]);
+        let _ = all_nan.quantile(0.5);
+        for pos in [0usize, 1, 31, 62, 63] {
+            let mut v = [0.0f64; 64];
+            v[pos] = f64::NAN;
+            let _ = Data::new(v).quantile(0.5);
+        }
+    }
+
     #[test]
     fn test_median_robust_on_infinities() {
         let data3 = [2.0, f64::NEG_INFINITY, f64::INFINITY];