Rust中如何获取Vec排序对应索引（等效numpy的argsort功能）

我们可以直接通过Rust标准库实现获取Vec排序后对应索引序列的功能，效果和numpy的argsort()完全一致，无需引入任何第三方依赖。
核心实现逻辑为：先生成存储原始数组所有下标的向量，再根据原始数组中对应下标的值对下标向量进行排序。

基础实现（适用于实现Ord trait的类型，如整数、字符串等） #

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

使用示例 #

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

扩展场景 #

1. 降序排序获取索引 #

如果需要按值从大到小的顺序生成索引，可以借助std::cmp::Reverse包装值即可：

use std::cmp::Reverse;

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

2. 处理浮点数排序 #

由于浮点数f32/f64未实现Ord trait（存在NaN的特殊情况），如果你能确认输入的浮点数中不存在NaN，可以按如下方式实现：

fn argsort_f64(arr: &[f64]) -> Vec<usize> {
    let mut indices: Vec<usize> = (0..arr.len()).collect();
    indices.sort_by(|&a, &b| arr[a].partial_cmp(&arr[b]).unwrap());
    indices
}

内容的提问来源于stack exchange，提问作者Michael Hall