How do I partially sort a Vec or slice?

Question

I need to get the top N items from a Vec which is quite large in production. Currently I do it like this inefficient way:

let mut v = vec![6, 4, 3, 7, 2, 1, 5];
v.sort_unstable();
v = v[0..3].to_vec();

In C++, I'd use std::partial_sort, but I can't find an equivalent in the Rust docs.

Am I just overlooking it, or does it not exist (yet)?

Answer 1

The standard library doesn't contain this functionality, but it looks like the lazysort crate is exactly what you need:

So what's the point of lazy sorting? As per the linked blog post, they're useful when you do not need or intend to need every value; for example you may only need the first 1,000 ordered values from a larger set.

#![feature(test)]

extern crate lazysort;
extern crate rand;
extern crate test;

use std::cmp::Ordering;

trait SortLazy<T> {
    fn sort_lazy<F>(&mut self, cmp: F, n: usize)
    where
        F: Fn(&T, &T) -> Ordering;
    unsafe fn sort_lazy_fast<F>(&mut self, cmp: F, n: usize)
    where
        F: Fn(&T, &T) -> Ordering;
}

impl<T> SortLazy<T> for [T] {
    fn sort_lazy<F>(&mut self, cmp: F, n: usize)
    where
        F: Fn(&T, &T) -> Ordering,
    {
        fn sort_lazy<F, T>(data: &mut [T], accu: &mut usize, cmp: &F, n: usize)
        where
            F: Fn(&T, &T) -> Ordering,
        {
            if !data.is_empty() && *accu < n {
                let mut pivot = 1;
                let mut lower = 0;
                let mut upper = data.len();
                while pivot < upper {
                    match cmp(&data[pivot], &data[lower]) {
                        Ordering::Less => {
                            data.swap(pivot, lower);
                            lower += 1;
                            pivot += 1;
                        }
                        Ordering::Greater => {
                            upper -= 1;
                            data.swap(pivot, upper);
                        }
                        Ordering::Equal => pivot += 1,
                    }
                }
                sort_lazy(&mut data[..lower], accu, cmp, n);
                sort_lazy(&mut data[upper..], accu, cmp, n);
            } else {
                *accu += 1;
            }
        }
        sort_lazy(self, &mut 0, &cmp, n);
    }

    unsafe fn sort_lazy_fast<F>(&mut self, cmp: F, n: usize)
    where
        F: Fn(&T, &T) -> Ordering,
    {
        fn sort_lazy<F, T>(data: &mut [T], accu: &mut usize, cmp: &F, n: usize)
        where
            F: Fn(&T, &T) -> Ordering,
        {
            if !data.is_empty() && *accu < n {
                unsafe {
                    use std::mem::swap;
                    let mut pivot = 1;
                    let mut lower = 0;
                    let mut upper = data.len();
                    while pivot < upper {
                        match cmp(data.get_unchecked(pivot), data.get_unchecked(lower)) {
                            Ordering::Less => {
                                swap(
                                    &mut *(data.get_unchecked_mut(pivot) as *mut T),
                                    &mut *(data.get_unchecked_mut(lower) as *mut T),
                                );
                                lower += 1;
                                pivot += 1;
                            }
                            Ordering::Greater => {
                                upper -= 1;
                                swap(
                                    &mut *(data.get_unchecked_mut(pivot) as *mut T),
                                    &mut *(data.get_unchecked_mut(upper) as *mut T),
                                );
                            }
                            Ordering::Equal => pivot += 1,
                        }
                    }
                    sort_lazy(&mut data[..lower], accu, cmp, n);
                    sort_lazy(&mut data[upper..], accu, cmp, n);
                }
            } else {
                *accu += 1;
            }
        }
        sort_lazy(self, &mut 0, &cmp, n);
    }
}

#[cfg(test)]
mod tests {
    use test::Bencher;

    use lazysort::Sorted;
    use std::collections::BinaryHeap;
    use SortLazy;

    use rand::{thread_rng, Rng};

    const SIZE_VEC: usize = 100_000;
    const N: usize = 42;

    #[bench]
    fn sort(b: &mut Bencher) {
        b.iter(|| {
            let mut rng = thread_rng();
            let mut v: Vec<i32> = std::iter::repeat_with(|| rng.gen())
                .take(SIZE_VEC)
                .collect();
            v.sort_unstable();
        })
    }

    #[bench]
    fn lazysort(b: &mut Bencher) {
        b.iter(|| {
            let mut rng = thread_rng();
            let v: Vec<i32> = std::iter::repeat_with(|| rng.gen())
                .take(SIZE_VEC)
                .collect();
            let _: Vec<_> = v.iter().sorted().take(N).collect();
        })
    }

    #[bench]
    fn lazysort_in_place(b: &mut Bencher) {
        b.iter(|| {
            let mut rng = thread_rng();
            let mut v: Vec<i32> = std::iter::repeat_with(|| rng.gen())
                .take(SIZE_VEC)
                .collect();
            v.sort_lazy(i32::cmp, N);
        })
    }

    #[bench]
    fn lazysort_in_place_fast(b: &mut Bencher) {
        b.iter(|| {
            let mut rng = thread_rng();
            let mut v: Vec<i32> = std::iter::repeat_with(|| rng.gen())
                .take(SIZE_VEC)
                .collect();
            unsafe { v.sort_lazy_fast(i32::cmp, N) };
        })
    }

    #[bench]
    fn binaryheap(b: &mut Bencher) {
        b.iter(|| {
            let mut rng = thread_rng();
            let v: Vec<i32> = std::iter::repeat_with(|| rng.gen())
                .take(SIZE_VEC)
                .collect();

            let mut iter = v.iter();
            let mut heap: BinaryHeap<_> = iter.by_ref().take(N).collect();
            for i in iter {
                heap.push(i);
                heap.pop();
            }
            let _ = heap.into_sorted_vec();
        })
    }
}

running 5 tests
test tests::binaryheap             ... bench:   3,283,938 ns/iter (+/- 413,805)
test tests::lazysort               ... bench:   1,669,229 ns/iter (+/- 505,528)
test tests::lazysort_in_place      ... bench:   1,781,007 ns/iter (+/- 443,472)
test tests::lazysort_in_place_fast ... bench:   1,652,103 ns/iter (+/- 691,847)
test tests::sort                   ... bench:   5,600,513 ns/iter (+/- 711,927)

test result: ok. 0 passed; 0 failed; 0 ignored; 5 measured; 0 filtered out

This code allows us to see that lazysort is faster than the solution with BinaryHeap. We can also see that BinaryHeap solution gets worse when N increases.

The problem with lazysort is that it creates a second Vec<_>. A "better" solution would be to implement the partial sort in-place. I provided an example of such an implementation.

Keep in mind that all these solutions come with overhead. When N is about SIZE_VEC / 3, the classic sort wins.

You could submit an RFC/issue to ask about adding this feature to the standard library.