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}