diff --git a/src/impl_1d.rs b/src/impl_1d.rs index b9bdb61f2..e69162259 100644 --- a/src/impl_1d.rs +++ b/src/impl_1d.rs @@ -27,13 +27,18 @@ impl ArrayBase } } - /// Return the element that would occupy the `i`-th position if the array - /// were sorted in increasing order. + /// Return the element that would occupy the `i`-th position if + /// the array were sorted in increasing order. /// /// The array is shuffled **in place** to retrieve the desired element: /// no copy of the array is allocated. - /// No assumptions should be made on the ordering of elements - /// after this computation. + /// After the shuffling, all elements with an index smaller than `i` + /// are smaller than the desired element, while all elements with + /// an index greater or equal than `i` are greater than or equal + /// to the desired element. + /// + /// No other assumptions should be made on the ordering of the + /// elements after this computation. /// /// Complexity ([quickselect](https://en.wikipedia.org/wiki/Quickselect)): /// - average case: O(`n`); diff --git a/src/numeric/impl_numeric.rs b/src/numeric/impl_numeric.rs index 9664c49ef..6ad427fcb 100644 --- a/src/numeric/impl_numeric.rs +++ b/src/numeric/impl_numeric.rs @@ -19,13 +19,69 @@ use { Zip, }; -/// Used to choose the interpolation strategy in [`percentile_axis_mut`]. +/// Used to provide an interpolation strategy to [`percentile_axis_mut`]. /// /// [`percentile_axis_mut`]: struct.ArrayBase.html#method.percentile_axis_mut -pub enum InterpolationStrategy { - Lower, - Nearest, - Higher, +pub trait Interpolate { + fn float_percentile_index(q: f64, len: usize) -> f64 { + ((len - 1) as f64) * q + } + + fn lower_index(q: f64, len: usize) -> usize { + Self::float_percentile_index(q, len).floor() as usize + } + + fn upper_index(q: f64, len: usize) -> usize { + Self::float_percentile_index(q, len).ceil() as usize + } + fn needs_lower(q: f64, len: usize) -> bool; + fn needs_upper(q: f64, len: usize) -> bool; + fn interpolate(lower: Option, upper: Option, q: f64, len: usize) -> T; +} + +pub struct Upper; +pub struct Lower; +pub struct Nearest; + +impl Interpolate for Upper { + fn needs_lower(_q: f64, _len: usize) -> bool { + false + } + fn needs_upper(_q: f64, _len: usize) -> bool { + true + } + fn interpolate(_lower: Option, upper: Option, _q: f64, _len: usize) -> T { + upper.unwrap() + } +} + +impl Interpolate for Lower { + fn needs_lower(_q: f64, _len: usize) -> bool { + true + } + fn needs_upper(_q: f64, _len: usize) -> bool { + false + } + fn interpolate(lower: Option, _upper: Option, _q: f64, _len: usize) -> T { + lower.unwrap() + } +} + +impl Interpolate for Nearest { + fn needs_lower(q: f64, len: usize) -> bool { + let lower = >::lower_index(q, len); + ((lower as f64) - >::float_percentile_index(q, len)) <= 0. + } + fn needs_upper(q: f64, len: usize) -> bool { + !>::needs_lower(q, len) + } + fn interpolate(lower: Option, upper: Option, q: f64, len: usize) -> T { + if >::needs_lower(q, len) { + lower.unwrap() + } else { + upper.unwrap() + } + } } /// Numerical methods for arrays. @@ -154,19 +210,33 @@ impl ArrayBase /// /// **Panics** if `axis` is out of bounds or if `q` is not between /// `0.` and `1.` (inclusive). - pub fn percentile_axis_mut(&mut self, axis: Axis, q: f64, interpolation_strategy: InterpolationStrategy) -> Array + pub fn percentile_axis_mut(&mut self, axis: Axis, q: f64) -> Array where D: RemoveAxis, A: Ord + Clone + Zero, S: DataMut, + I: Interpolate>, { assert!((0. <= q) && (q <= 1.)); - let float_percentile_index = ((self.len_of(axis) - 1) as f64) * q; - let percentile_index = match interpolation_strategy { - InterpolationStrategy::Lower => float_percentile_index.floor() as usize, - InterpolationStrategy::Nearest => float_percentile_index.round() as usize, - InterpolationStrategy::Higher => float_percentile_index.ceil() as usize, + let mut lower = None; + let mut upper = None; + let axis_len = self.len_of(axis); + if I::needs_lower(q, axis_len) { + lower = Some( + self.map_axis_mut( + axis, + |mut x| x.sorted_get_mut(I::lower_index(q, axis_len)) + ) + ); + }; + if I::needs_upper(q, axis_len) { + upper = Some( + self.map_axis_mut( + axis, + |mut x| x.sorted_get_mut(I::upper_index(q, axis_len)) + ) + ); }; - self.map_axis_mut(axis, |mut x| x.sorted_get_mut(percentile_index)) + I::interpolate(lower, upper, q, axis_len) } /// Return variance along `axis`. diff --git a/src/numeric/mod.rs b/src/numeric/mod.rs index 88ff33649..30024d902 100644 --- a/src/numeric/mod.rs +++ b/src/numeric/mod.rs @@ -1,3 +1,3 @@ +pub use self::impl_numeric::*; -pub mod impl_numeric; -pub use self::impl_numeric::InterpolationStrategy; +mod impl_numeric; diff --git a/tests/array.rs b/tests/array.rs index 1369c1be6..75a2610fa 100644 --- a/tests/array.rs +++ b/tests/array.rs @@ -8,7 +8,7 @@ extern crate itertools; use ndarray::{Slice, SliceInfo, SliceOrIndex}; use ndarray::prelude::*; -use ndarray::numeric::InterpolationStrategy; +use ndarray::numeric::{Lower}; use ndarray::{ rcarr2, arr3, @@ -1759,7 +1759,7 @@ fn test_percentile_axis_mut_with_odd_axis_length() { [3, 5, 6, 12] ] ); - let p = a.percentile_axis_mut(Axis(0), 0.5, InterpolationStrategy::Lower); + let p = a.percentile_axis_mut::(Axis(0), 0.5); assert!(p == a.subview(Axis(0), 1)); } @@ -1773,20 +1773,20 @@ fn test_percentile_axis_mut_with_even_axis_length() { [4, 6, 7, 13] ] ); - let q = b.percentile_axis_mut(Axis(0), 0.5, InterpolationStrategy::Lower); + let q = b.percentile_axis_mut::(Axis(0), 0.5); assert!(q == b.subview(Axis(0), 1)); } #[test] fn test_percentile_axis_mut_to_get_minimum() { let mut b = arr2(&[[1, 3, 22, 10]]); - let q = b.percentile_axis_mut(Axis(1), 0., InterpolationStrategy::Lower); + let q = b.percentile_axis_mut::(Axis(1), 0.); assert!(q == arr1(&[1])); } #[test] fn test_percentile_axis_mut_to_get_maximum() { let mut b = arr1(&[1, 3, 22, 10]); - let q = b.percentile_axis_mut(Axis(0), 1., InterpolationStrategy::Lower); + let q = b.percentile_axis_mut::(Axis(0), 1.); assert!(q == arr0(22)); }