Home > Enterprise >  How to get the indices that would sort a vector in Rust?
How to get the indices that would sort a vector in Rust?

Time:10-29

I want to get a vector of the indices that would sort a vector in Rust. Effectively, I want argsort from numpy.

For example

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

CodePudding user response:

Not sure if there's something pre-made for this, but its simple enough to implement yourself with .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
}

See it working on the playground.

  • Related