How to get the indicies that would sort a vector of floats in rust?-CodePudding

Sorting a vector of integers is straight forward as demonstrated here. However, a vector of floats is more complicated due to potential NaN's and floating point operations.

I would like to "combine" the two methods below to get the indicies that would sort a vector of floats, without sorting the actual input vector.

// returns the indices that would sort a vector of ints
fn argsort<T: Ord>(data: &[T]) -> Vec<usize> {
    let mut indices = (0..data.len()).collect::<Vec<_>>();
    indices.sort_by_key(|&i| &data[i]);
    indices
}

and

use std::cmp::Ordering;


// returns a sorted vector of floats that may contain NaNs, with NaNs at the end
fn sort(arr: &Vec<f64>) -> Vec<f64> {
    let mut out = arr.clone();
    out.sort_by(|&a, &b| {
        match (a.is_nan(), b.is_nan()) {
            (true, true) => Ordering::Equal,
            (true, false) => Ordering::Greater,
            (false, true) => Ordering::Less,
            (false, false) => a.partial_cmp(&b).unwrap(),
        }
    });
    return out;
}

How can I return a vector of usize indices that would sort an input vector of floats containing NaNs?

CodePudding user response：

i'm not sure if i understood your question. It seems that your second snippet almost answers the question. Is this the answer you were looking for?

   fn sort(arr: &Vec<f64>) -> Vec<usize> {
    let mut out = (0..arr.len()).collect::<Vec<usize>>();
    out.sort_by(|&a_idx, &b_idx| {
        let a = arr[a_idx];
        let b = arr[b_idx];
        match (a.is_nan(), b.is_nan()) {
            (true, true) => Ordering::Equal,
            (true, false) => Ordering::Greater,
            (false, true) => Ordering::Less,
            (false, false) => a.partial_cmp(&b).unwrap(),
        }
    });
    out
}

CodePudding user response：

// construct a tuple vector from the floats, the tuple being (index,float) 
let mut with_indices = arr.clone().into_iter().enumerate()
       .map(|index,value| (index,value)).collect::<Vec<(usize,f64)>>();

// sort the tuple vector by float
with_indices.sort_by(|&a, &b| {
    match (a.1.is_nan(), b.1.is_nan()) {
        (true, true) => Ordering::Equal,
        (true, false) => Ordering::Greater,
        (false, true) => Ordering::Less,
        (false, false) => a.1.partial_cmp(&b.1).unwrap(),
    }
});

// extract the sorted indices from the sorted tuple vector
let stored_indices = with_indices.into_iter().map(|(index,value)| index)
             .collect::<Vec<usize>>();