如何在Rust中得到排序向量的索引?

9rnv2umw  于 2022-12-19  发布在  其他
关注(0)|答案(2)|浏览(141)

我想得到一个索引的向量,它可以对Rust中的向量进行排序,实际上,我想得到numpy中的argsort
例如

let v = vec![1, 7, 4, 2];
let i = argsort(&v);
assert_eq!(i, &[0, 3, 2, 1]);
yeotifhr

yeotifhr1#

不确定是否有预先准备的东西,但它足够简单,可以用.sort_by_key()实现自己:

pub fn argsort<T: Ord>(data: &[T]) -> Vec<usize> {
    let mut indices = (0..data.len()).collect::<Vec<_>>();
    indices.sort_by_key(|&i| &data[i]);
    indices
}

请看它在playground上的工作情况。

ezykj2lf

ezykj2lf2#

一年多来,我一直在使用@kmdreko提供的答案,效果很好,但我有一些应用程序,其中这个函数对性能至关重要,所以我进行了基准测试,并提出了两种替代的、更快的方法。
第一种方法通过使用enumerate完全避免了创建新的索引向量。

pub fn argsort_enumerate<T: Ord>(data: &[T]) -> Vec<usize> {
    let mut indices: Vec<_> = data.iter().enumerate().collect();
    indices.sort_by_key(|&(_, v)| v);
    // Extract the indices from the sorted vector
    indices.into_iter().map(|(i, _)| i).collect()
}

第二个版本是使用rayon的并行版本,它允许并行排序,如果数据很大并且排序操作是计算密集型的(我的应用程序),这可以提高性能。

use rayon::slice::ParallelSliceMut;

pub fn argsort_par<T: Ord + Sync>(data: &[T]) -> Vec<usize> {
    let mut indices = (0..data.len()).collect::<Vec<_>>();
    indices.par_sort_unstable_by_key(|&i| &data[i]);
    indices
}

注意,需要添加一个Sync绑定。
我使用Criterion对这三个实现(来自@kmdreko的实现称为argsort_simple)进行了基准测试,测试对象是从0到2000000之间的随机均匀分布中采样的10,000个u64向量

argsort_simple          time:   [738.07 µs 740.26 µs 742.78 µs]

Found 13 outliers among 100 measurements (13.00%)
  7 (7.00%) high mild
  6 (6.00%) high severe

argsort_par             time:   [162.10 µs 163.72 µs 165.72 µs]

Found 10 outliers among 100 measurements (10.00%)
  5 (5.00%) high mild
  5 (5.00%) high severe

argsort_enumerate       time:   [594.97 µs 600.31 µs 605.96 µs]

Found 4 outliers among 100 measurements (4.00%)
  1 (1.00%) high mild
  3 (3.00%) high severe

相关问题