3

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]);
Michael Hall
  • 2,141
  • 1
  • 18
  • 33

1 Answers1

14

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.

kmdreko
  • 25,172
  • 5
  • 32
  • 61
  • 1
    What about for a vector of floats? – Katya Feb 20 '22 at 18:13
  • 2
    @Katya Floats do not have a total order. You would need `sort_by_key` to return a wrapper type or use [`sort_by()`](https://doc.rust-lang.org/std/primitive.slice.html#method.sort_by) with a function that can deal with comparing `NaN`. See answers [here](https://stackoverflow.com/q/28247990) or [here](https://stackoverflow.com/q/40408293) for ideas. – kmdreko Feb 20 '22 at 18:27
  • Is there a reason a built in or library packaged function such as `argsort()` for float data types is not available? I am coming over from Python and things seem a bit convoluted sometimes in Rust (I am sure I will learn to appreciate it immensely with more use). – Katya Feb 21 '22 at 11:57