select.rs - source
Description: Source of the Rust file `library/core/src/slice/sort/select.rs`.

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