diff --git a/src/sort.rs b/src/sort.rs index e38e205..e0c2002 100644 --- a/src/sort.rs +++ b/src/sort.rs @@ -96,25 +96,31 @@ pub trait Sort1dExt { } impl Sort1dExt for ArrayRef { - fn get_from_sorted_mut(&mut self, i: usize) -> A + fn get_from_sorted_mut(&mut self, mut i: usize) -> A where A: Ord + Clone, { - let n = self.len(); - if n == 1 { - self[0].clone() - } else { + // We narrow the view in place instead of recursing, so that the stack + // depth stays constant regardless of the number of partition steps. + // This matters for pathological inputs (e.g. large constant arrays), + // where each partition step only peels off a single element and the + // recursive formulation would use O(n) stack frames. + let mut array = self.view_mut(); + loop { + let n = array.len(); + if n == 1 { + return array[0].clone(); + } let mut rng = thread_rng(); let pivot_index = rng.gen_range(0..n); - let partition_index = self.partition_mut(pivot_index); + let partition_index = array.partition_mut(pivot_index); if i < partition_index { - self.slice_axis_mut(Axis(0), Slice::from(..partition_index)) - .get_from_sorted_mut(i) + array.slice_axis_inplace(Axis(0), Slice::from(..partition_index)); } else if i == partition_index { - self[i].clone() + return array[i].clone(); } else { - self.slice_axis_mut(Axis(0), Slice::from(partition_index + 1..)) - .get_from_sorted_mut(i - (partition_index + 1)) + array.slice_axis_inplace(Axis(0), Slice::from(partition_index + 1..)); + i -= partition_index + 1; } } } @@ -211,73 +217,86 @@ where /// initial element values are ignored. fn _get_many_from_sorted_mut_unchecked( mut array: ArrayViewMut1<'_, A>, - indexes: &mut [usize], - values: &mut [A], + mut indexes: &mut [usize], + mut values: &mut [A], ) where A: Ord + Clone, { - let n = array.len(); - debug_assert!(n >= indexes.len()); // because indexes must be unique and in-bounds - debug_assert_eq!(indexes.len(), values.len()); + // After each partition step we recurse into the sub-problem with *fewer* + // array elements and iterate (in this `loop`) on the larger one. This is + // the standard technique for bounding quickselect's stack depth to + // O(log n): without it, pathological inputs such as large constant arrays + // (where each partition step only peels off a single element) would use + // O(n) stack frames and overflow the stack. + loop { + let n = array.len(); + debug_assert!(n >= indexes.len()); // because indexes must be unique and in-bounds + debug_assert_eq!(indexes.len(), values.len()); - if indexes.is_empty() { - // Nothing to do in this case. - return; - } + if indexes.is_empty() { + // Nothing to do in this case. + return; + } - // At this point, `n >= 1` since `indexes.len() >= 1`. - if n == 1 { - // We can only reach this point if `indexes.len() == 1`, so we only - // need to assign the single value, and then we're done. - debug_assert_eq!(indexes.len(), 1); - values[0] = array[0].clone(); - return; - } + // At this point, `n >= 1` since `indexes.len() >= 1`. + if n == 1 { + // We can only reach this point if `indexes.len() == 1`, so we only + // need to assign the single value, and then we're done. + debug_assert_eq!(indexes.len(), 1); + values[0] = array[0].clone(); + return; + } + + // We pick a random pivot index: the corresponding element is the pivot value + let mut rng = thread_rng(); + let pivot_index = rng.gen_range(0..n); - // We pick a random pivot index: the corresponding element is the pivot value - let mut rng = thread_rng(); - let pivot_index = rng.gen_range(0..n); + // We partition the array with respect to the pivot value. + // The pivot value moves to `array_partition_index`. + // Elements strictly smaller than the pivot value have indexes < `array_partition_index`. + // Elements greater or equal to the pivot value have indexes > `array_partition_index`. + let array_partition_index = array.partition_mut(pivot_index); - // We partition the array with respect to the pivot value. - // The pivot value moves to `array_partition_index`. - // Elements strictly smaller than the pivot value have indexes < `array_partition_index`. - // Elements greater or equal to the pivot value have indexes > `array_partition_index`. - let array_partition_index = array.partition_mut(pivot_index); + // We use a divide-and-conquer strategy, splitting the indexes we are + // searching for (`indexes`) and the corresponding portions of the output + // slice (`values`) into pieces with respect to `array_partition_index`. + let (found_exact, index_split) = match indexes.binary_search(&array_partition_index) { + Ok(index) => (true, index), + Err(index) => (false, index), + }; + let (smaller_indexes, other_indexes) = indexes.split_at_mut(index_split); + let (smaller_values, other_values) = values.split_at_mut(index_split); + let (bigger_indexes, bigger_values) = if found_exact { + other_values[0] = array[array_partition_index].clone(); // Write exactly found value. + (&mut other_indexes[1..], &mut other_values[1..]) + } else { + (other_indexes, other_values) + }; - // We use a divide-and-conquer strategy, splitting the indexes we are - // searching for (`indexes`) and the corresponding portions of the output - // slice (`values`) into pieces with respect to `array_partition_index`. - let (found_exact, index_split) = match indexes.binary_search(&array_partition_index) { - Ok(index) => (true, index), - Err(index) => (false, index), - }; - let (smaller_indexes, other_indexes) = indexes.split_at_mut(index_split); - let (smaller_values, other_values) = values.split_at_mut(index_split); - let (bigger_indexes, bigger_values) = if found_exact { - other_values[0] = array[array_partition_index].clone(); // Write exactly found value. - (&mut other_indexes[1..], &mut other_values[1..]) - } else { - (other_indexes, other_values) - }; + // The indexes to the right of the pivot need to be shifted by the length + // of the removed portion (the smaller part plus the pivot itself). + bigger_indexes + .iter_mut() + .for_each(|x| *x -= array_partition_index + 1); - // We search recursively for the values corresponding to strictly smaller - // indexes to the left of `partition_index`. - _get_many_from_sorted_mut_unchecked( - array.slice_axis_mut(Axis(0), Slice::from(..array_partition_index)), - smaller_indexes, - smaller_values, - ); + // We split the array into the two disjoint sub-views, dropping the pivot + // (which is already in its final sorted position and, if requested, has + // already been written to `bigger_values`/`other_values` above). + let (smaller_array, rest) = array.split_at(Axis(0), array_partition_index); + let bigger_array = rest.slice_axis_move(Axis(0), Slice::from(1..)); - // We search recursively for the values corresponding to strictly bigger - // indexes to the right of `partition_index`. Since only the right portion - // of the array is passed in, the indexes need to be shifted by length of - // the removed portion. - bigger_indexes - .iter_mut() - .for_each(|x| *x -= array_partition_index + 1); - _get_many_from_sorted_mut_unchecked( - array.slice_axis_mut(Axis(0), Slice::from(array_partition_index + 1..)), - bigger_indexes, - bigger_values, - ); + // Recurse into the sub-problem with fewer array elements and continue + // iterating on the larger one, so the recursion depth stays O(log n). + if smaller_array.len() <= bigger_array.len() { + _get_many_from_sorted_mut_unchecked(smaller_array, smaller_indexes, smaller_values); + array = bigger_array; + indexes = bigger_indexes; + values = bigger_values; + } else { + _get_many_from_sorted_mut_unchecked(bigger_array, bigger_indexes, bigger_values); + array = smaller_array; + indexes = smaller_indexes; + values = smaller_values; + } + } } diff --git a/tests/quantile.rs b/tests/quantile.rs index 9d58071..a103df1 100644 --- a/tests/quantile.rs +++ b/tests/quantile.rs @@ -277,6 +277,14 @@ fn test_midpoint_overflow() { assert_eq!(median, expected_median); } +#[test] +fn quantile_mut_on_large_constant_array_does_not_overflow_stack() { + // Regression test for https://github.com/rust-ndarray/ndarray-stats/issues/86 + let mut array: Array1 = Array1::ones(100_000); + let median = array.quantile_mut(n64(0.5), &Linear).unwrap(); + assert_eq!(median, n64(1.)); +} + #[quickcheck] fn test_quantiles_mut(xs: Vec) -> bool { let v = Array::from(xs.clone());