core/slice/sort/
select.rs

1//! This module contains the implementation for `slice::select_nth_unstable`.
2//! It uses an introselect algorithm based on ipnsort by Lukas Bergdoll and Orson Peters,
3//! published at: <https://github.com/Voultapher/sort-research-rs/tree/main/ipnsort>
4//!
5//! The fallback algorithm used for introselect is Median of Medians using Tukey's Ninther
6//! for pivot selection. Using this as a fallback ensures O(n) worst case running time with
7//! better performance than one would get using heapsort as fallback.
8
9use crate::mem::{self, SizedTypeProperties};
10#[cfg(not(feature = "optimize_for_size"))]
11use crate::slice::sort::shared::pivot::choose_pivot;
12use crate::slice::sort::shared::smallsort::insertion_sort_shift_left;
13use crate::slice::sort::unstable::quicksort::partition;
14
15/// Reorders the slice such that the element at `index` is at its final sorted position.
16pub(crate) fn partition_at_index<T, F>(
17    v: &mut [T],
18    index: usize,
19    mut is_less: F,
20) -> (&mut [T], &mut T, &mut [T])
21where
22    F: FnMut(&T, &T) -> bool,
23{
24    let len = v.len();
25
26    // Puts a lower limit of 1 on `len`.
27    if index >= len {
28        panic!("partition_at_index index {} greater than length of slice {}", index, len);
29    }
30
31    if T::IS_ZST {
32        // Sorting has no meaningful behavior on zero-sized types. Do nothing.
33    } else if index == len - 1 {
34        // Find max element and place it in the last position of the array. We're free to use
35        // `unwrap()` here because we checked that `v` is not empty.
36        let max_idx = max_index(v, &mut is_less).unwrap();
37        v.swap(max_idx, index);
38    } else if index == 0 {
39        // Find min element and place it in the first position of the array. We're free to use
40        // `unwrap()` here because we checked that `v` is not empty.
41        let min_idx = min_index(v, &mut is_less).unwrap();
42        v.swap(min_idx, index);
43    } else {
44        cfg_if! {
45            if #[cfg(feature = "optimize_for_size")] {
46                median_of_medians(v, &mut is_less, index);
47            } else {
48                partition_at_index_loop(v, index, None, &mut is_less);
49            }
50        }
51    }
52
53    let (left, right) = v.split_at_mut(index);
54    let (pivot, right) = right.split_at_mut(1);
55    let pivot = &mut pivot[0];
56    (left, pivot, right)
57}
58
59// For small sub-slices it's faster to use a dedicated small-sort, but because it is only called at
60// most once, it doesn't make sense to use something more sophisticated than insertion-sort.
61const INSERTION_SORT_THRESHOLD: usize = 16;
62
63#[cfg(not(feature = "optimize_for_size"))]
64fn partition_at_index_loop<'a, T, F>(
65    mut v: &'a mut [T],
66    mut index: usize,
67    mut ancestor_pivot: Option<&'a T>,
68    is_less: &mut F,
69) where
70    F: FnMut(&T, &T) -> bool,
71{
72    // Limit the amount of iterations and fall back to fast deterministic selection to ensure O(n)
73    // worst case running time. This limit needs to be constant, because using `ilog2(len)` like in
74    // `sort` would result in O(n log n) time complexity. The exact value of the limit is chosen
75    // somewhat arbitrarily, but for most inputs bad pivot selections should be relatively rare, so
76    // the limit is reached for sub-slices len / (2^limit or less). Which makes the remaining work
77    // with the fallback minimal in relative terms.
78    let mut limit = 16;
79
80    loop {
81        if v.len() <= INSERTION_SORT_THRESHOLD {
82            if v.len() >= 2 {
83                insertion_sort_shift_left(v, 1, is_less);
84            }
85            return;
86        }
87
88        if limit == 0 {
89            median_of_medians(v, is_less, index);
90            return;
91        }
92
93        limit -= 1;
94
95        // Choose a pivot
96        let pivot_pos = choose_pivot(v, is_less);
97
98        // If the chosen pivot is equal to the predecessor, then it's the smallest element in the
99        // slice. Partition the slice into elements equal to and elements greater than the pivot.
100        // This case is usually hit when the slice contains many duplicate elements.
101        if let Some(p) = ancestor_pivot {
102            // SAFETY: choose_pivot promises to return a valid pivot position.
103            let pivot = unsafe { v.get_unchecked(pivot_pos) };
104
105            if !is_less(p, pivot) {
106                let num_lt = partition(v, pivot_pos, &mut |a, b| !is_less(b, a));
107
108                // Continue sorting elements greater than the pivot. We know that `mid` contains
109                // the pivot. So we can continue after `mid`.
110                let mid = num_lt + 1;
111
112                // If we've passed our index, then we're good.
113                if mid > index {
114                    return;
115                }
116
117                v = &mut v[mid..];
118                index = index - mid;
119                ancestor_pivot = None;
120                continue;
121            }
122        }
123
124        let mid = partition(v, pivot_pos, is_less);
125
126        // Split the slice into `left`, `pivot`, and `right`.
127        let (left, right) = v.split_at_mut(mid);
128        let (pivot, right) = right.split_at_mut(1);
129        let pivot = &pivot[0];
130
131        if mid < index {
132            v = right;
133            index = index - mid - 1;
134            ancestor_pivot = Some(pivot);
135        } else if mid > index {
136            v = left;
137        } else {
138            // If mid == index, then we're done, since partition() guaranteed that all elements
139            // after mid are greater than or equal to mid.
140            return;
141        }
142    }
143}
144
145/// Helper function that returns the index of the minimum element in the slice using the given
146/// comparator function
147fn min_index<T, F: FnMut(&T, &T) -> bool>(slice: &[T], is_less: &mut F) -> Option<usize> {
148    slice
149        .iter()
150        .enumerate()
151        .reduce(|acc, t| if is_less(t.1, acc.1) { t } else { acc })
152        .map(|(i, _)| i)
153}
154
155/// Helper function that returns the index of the maximum element in the slice using the given
156/// comparator function
157fn max_index<T, F: FnMut(&T, &T) -> bool>(slice: &[T], is_less: &mut F) -> Option<usize> {
158    slice
159        .iter()
160        .enumerate()
161        .reduce(|acc, t| if is_less(acc.1, t.1) { t } else { acc })
162        .map(|(i, _)| i)
163}
164
165/// Selection algorithm to select the k-th element from the slice in guaranteed O(n) time.
166/// This is essentially a quickselect that uses Tukey's Ninther for pivot selection
167fn median_of_medians<T, F: FnMut(&T, &T) -> bool>(mut v: &mut [T], is_less: &mut F, mut k: usize) {
168    // Since this function isn't public, it should never be called with an out-of-bounds index.
169    debug_assert!(k < v.len());
170
171    // If T is as ZST, `partition_at_index` will already return early.
172    debug_assert!(!T::IS_ZST);
173
174    // We now know that `k < v.len() <= isize::MAX`
175    loop {
176        if v.len() <= INSERTION_SORT_THRESHOLD {
177            if v.len() >= 2 {
178                insertion_sort_shift_left(v, 1, is_less);
179            }
180
181            return;
182        }
183
184        // `median_of_{minima,maxima}` can't handle the extreme cases of the first/last element,
185        // so we catch them here and just do a linear search.
186        if k == v.len() - 1 {
187            // Find max element and place it in the last position of the array. We're free to use
188            // `unwrap()` here because we know v must not be empty.
189            let max_idx = max_index(v, is_less).unwrap();
190            v.swap(max_idx, k);
191            return;
192        } else if k == 0 {
193            // Find min element and place it in the first position of the array. We're free to use
194            // `unwrap()` here because we know v must not be empty.
195            let min_idx = min_index(v, is_less).unwrap();
196            v.swap(min_idx, k);
197            return;
198        }
199
200        let p = median_of_ninthers(v, is_less);
201
202        if p == k {
203            return;
204        } else if p > k {
205            v = &mut v[..p];
206        } else {
207            // Since `p < k < v.len()`, `p + 1` doesn't overflow and is
208            // a valid index into the slice.
209            v = &mut v[p + 1..];
210            k -= p + 1;
211        }
212    }
213}
214
215// Optimized for when `k` lies somewhere in the middle of the slice. Selects a pivot
216// as close as possible to the median of the slice. For more details on how the algorithm
217// operates, refer to the paper <https://drops.dagstuhl.de/opus/volltexte/2017/7612/pdf/LIPIcs-SEA-2017-24.pdf>.
218fn median_of_ninthers<T, F: FnMut(&T, &T) -> bool>(v: &mut [T], is_less: &mut F) -> usize {
219    // use `saturating_mul` so the multiplication doesn't overflow on 16-bit platforms.
220    let frac = if v.len() <= 1024 {
221        v.len() / 12
222    } else if v.len() <= 128_usize.saturating_mul(1024) {
223        v.len() / 64
224    } else {
225        v.len() / 1024
226    };
227
228    let pivot = frac / 2;
229    let lo = v.len() / 2 - pivot;
230    let hi = frac + lo;
231    let gap = (v.len() - 9 * frac) / 4;
232    let mut a = lo - 4 * frac - gap;
233    let mut b = hi + gap;
234    for i in lo..hi {
235        ninther(v, is_less, a, i - frac, b, a + 1, i, b + 1, a + 2, i + frac, b + 2);
236        a += 3;
237        b += 3;
238    }
239
240    median_of_medians(&mut v[lo..lo + frac], is_less, pivot);
241
242    partition(v, lo + pivot, is_less)
243}
244
245/// Moves around the 9 elements at the indices a..i, such that
246/// `v[d]` contains the median of the 9 elements and the other
247/// elements are partitioned around it.
248fn ninther<T, F: FnMut(&T, &T) -> bool>(
249    v: &mut [T],
250    is_less: &mut F,
251    a: usize,
252    mut b: usize,
253    c: usize,
254    mut d: usize,
255    e: usize,
256    mut f: usize,
257    g: usize,
258    mut h: usize,
259    i: usize,
260) {
261    b = median_idx(v, is_less, a, b, c);
262    h = median_idx(v, is_less, g, h, i);
263    if is_less(&v[h], &v[b]) {
264        mem::swap(&mut b, &mut h);
265    }
266    if is_less(&v[f], &v[d]) {
267        mem::swap(&mut d, &mut f);
268    }
269    if is_less(&v[e], &v[d]) {
270        // do nothing
271    } else if is_less(&v[f], &v[e]) {
272        d = f;
273    } else {
274        if is_less(&v[e], &v[b]) {
275            v.swap(e, b);
276        } else if is_less(&v[h], &v[e]) {
277            v.swap(e, h);
278        }
279        return;
280    }
281    if is_less(&v[d], &v[b]) {
282        d = b;
283    } else if is_less(&v[h], &v[d]) {
284        d = h;
285    }
286
287    v.swap(d, e);
288}
289
290/// returns the index pointing to the median of the 3
291/// elements `v[a]`, `v[b]` and `v[c]`
292fn median_idx<T, F: FnMut(&T, &T) -> bool>(
293    v: &[T],
294    is_less: &mut F,
295    mut a: usize,
296    b: usize,
297    mut c: usize,
298) -> usize {
299    if is_less(&v[c], &v[a]) {
300        mem::swap(&mut a, &mut c);
301    }
302    if is_less(&v[c], &v[b]) {
303        return c;
304    }
305    if is_less(&v[b], &v[a]) {
306        return a;
307    }
308    b
309}