1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132

use std::{cmp, ptr};
use std::cmp::Ordering;

pub fn select<T, F>(array: &mut [T], k: usize, mut f: F)
    where F: FnMut(&T, &T) -> cmp::Ordering
{
    let r = array.len() - 1;
    select_(array, &mut f, 0, r, k)
}

const A: usize = 600;
const B: f32 = 0.5;

fn select_<T, F>(array: &mut [T], cmp: &mut F, mut left: usize, mut right: usize, k: usize)
    where F: FnMut(&T, &T) -> cmp::Ordering
{
    let array = array;
    while right > left {
        if right - left > A {
            let n = (right - left + 1) as f32;
            let i = (k - left + 1) as f32;
            let z = n.ln();
            let s = B * (z * (2.0 / 3.0)).exp();
            let sn = s / n;
            let sd = B * (z * s * (1.0 - sn)).sqrt() * (i - n * 0.5).signum();

            let isn = i * s / n;
            let inner = k as f32 - isn + sd;
            let new_left = cmp::max(left, inner as usize);
            let new_right = cmp::min(right, (inner + s) as usize);

            select_(array, cmp, new_left, new_right, k)
        }

        let mut i = left + 1;
        let mut j = right - 1;
        array.swap(left, k);
        let t_idx = if cmp(&array[left], &array[right]) != cmp::Ordering::Less {
            array.swap(left, right);
            right
        } else {
            left
        };

        // need to cancel the borrow (but the assertion above ensures this doesn't alias)
        let t = &array[t_idx] as *const _;
        unsafe {
            while cmp(&*array.get_unchecked(i), &*t) == Ordering::Less { i += 1 }
            while cmp(&*array.get_unchecked(j), &*t) == Ordering::Greater { j -= 1 }
        }

        if i < j {
            // i < j, and i and j move toward each other, so this
            // assertion ensures that all indexing here is in-bounds.
            assert!(j < array.len());

            // FIXME: this unsafe code *should* be unnecessary: the
            // assertions above mean that LLVM could theoretically
            // optimise out the bounds checks, but it doesn't seem to
            // at the moment (2015-04-25).
            unsafe {
                while i < j {
                    ptr::swap(array.get_unchecked_mut(i),
                              array.get_unchecked_mut(j));
                    i += 1;
                    j -= 1;
                    while cmp(&*array.get_unchecked(i), &*t) == Ordering::Less { i += 1 }
                    while cmp(&*array.get_unchecked(j), &*t) == Ordering::Greater { j -= 1 }
                }
            }
        }

        if left == t_idx {
            array.swap(left, j);
        } else {
            j += 1;
            array.swap(right, j);
        }
        if j <= k { left = j + 1 }
        if k <= j { right = j.saturating_sub(1); }
    }
}

#[cfg(test)]
mod tests {
    use super::select;
    use quickcheck::{self, TestResult};
    use rand::{XorShiftRng, Rng};

    #[test]
    fn qc() {
        fn run(mut x: Vec<i32>, k: usize) -> TestResult {
            if k >= x.len() {
                return TestResult::discard();
            }

            select(&mut x, k, Ord::cmp);
            let element = x[k];
            x.sort();
            TestResult::from_bool(element == x[k])
        }
        quickcheck::quickcheck(run as fn(Vec<i32>, usize) -> TestResult)
    }

    #[test]
    fn smoke() {
        for k in 0..4 {
            let mut array = [2, 3, 0, 1];
            select(&mut array, k, Ord::cmp);
            assert_eq!(array[k], k);
        }
    }
    #[test]
    fn huge() {
        let mut rng = XorShiftRng::new_unseeded();
        for _ in 0..20 {
            let length = rng.gen_range(2_000, 10_000);
            let v: Vec<_> = rng.gen_iter::<i32>().take(length).collect();
            for _ in 0..10 {
                let mut w = v.clone();
                select(&mut w, rng.gen_range(0, length), Ord::cmp);
            }
        }
    }
}

#[cfg(all(test, feature = "unstable"))]
mod benches {
    extern crate test;

    make_benches!(|m, mut v| super::select(&mut v, m, Ord::cmp));
}