performancesortingvectorrustpartial-sort

How do I partially sort a Vec or slice?


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)?


Solution

  • 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.