diff --git a/benches/order_statistics.rs b/benches/order_statistics.rs index d1771ff4..c103143a 100644 --- a/benches/order_statistics.rs +++ b/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 && 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 { + let mut rng = StdRng::seed_from_u64(0xB0A7); + match shape { + // shuffled, the ordinary case + "random" => { + let mut v: Vec = (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 = (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); diff --git a/src/statistics/slice_statistics.rs b/src/statistics/slice_statistics.rs index 96b83e7b..62840b25 100644 --- a/src/statistics/slice_statistics.rs +++ b/src/statistics/slice_statistics.rs @@ -64,8 +64,19 @@ impl + AsRef<[f64]>> Data { 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 + AsRef<[f64]>> Data { 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 { + 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 = (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];