如何对 Vec 或切片进行部分排序？

Question

我需要从生产中相当大的 Vec 中获取前 N 个项目。目前我这样做效率低下：

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

在 C++ 中，我会使用 std::partial_sort, but I can't find an equivalent in the Rust docs。

我只是忽略了它，还是它（还）不存在？

Answer 1

标准库不包含此功能，但看起来 lazysort crate 正是您所需要的：

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

这段代码让我们看到 lazysort 比 BinaryHeap 的解决方案更快。我们还可以看到，当 N 增加时，BinaryHeap 解决方案变得更糟。

lazysort 的问题是它创建了第二个 Vec<_>。 "better" 解决方案是就地实施部分排序。我提供了此类实现的示例。

请记住，所有这些解决方案都会带来开销。当N约SIZE_VEC / 3时，经典sort胜

您可以提交 RFC/issue 来询问是否将此功能添加到标准库中。

如何对 Vec 或切片进行部分排序？

How do I partially sort a Vec or slice?

sorting

performance

vector

rust

partial-sort