如何在Rust中高效按指定列排序ndarray::Array2？

要高效对ndarray::Array2按指定列排序，核心思路是利用ndarray的行迭代器结合Rust标准库的排序方法，直接操作数组行元素。以下是具体实现方案：

1. 依赖准备 #

确保Cargo.toml中引入ndarray依赖：

[dependencies]
ndarray = "0.15.6" # 使用最新稳定版即可

2. 升序排序实现 #

通过axis_iter_mut获取行的可变迭代器，使用sort_by_key提取指定列的值作为排序键，完成升序排序：

use ndarray::{Array2, Axis};

fn sort_array_by_col_asc(arr: &mut Array2<i32>, col_idx: usize) {
    // 检查列索引合法性，避免越界
    assert!(col_idx < arr.ncols(), "Column index out of bounds");
    
    // 稳定升序排序（保留相等元素原顺序）
    arr.axis_iter_mut(Axis(0))
        .sort_by_key(|row| row[col_idx]);
    
    // 若追求更高性能（无需稳定排序），可替换为不稳定排序：
    // arr.axis_iter_mut(Axis(0))
    //     .sort_unstable_by_key(|row| row[col_idx]);
}

fn main() {
    let mut arr = Array2::from_shape_vec((3, 3), vec![3, 1, 2, 1, 3, 2, 2, 1, 3]).unwrap();
    println!("原数组:\n{}", arr);
    
    sort_array_by_col_asc(&mut arr, 1);
    println!("按第2列升序排序后:\n{}", arr);
}

3. 降序排序实现 #

降序排序有两种简洁实现方式：

方法一：使用Reverse包装排序键（推荐） #

借助std::cmp::Reverse反转排序逻辑，代码更简洁：

use ndarray::{Array2, Axis};
use std::cmp::Reverse;

fn sort_array_by_col_desc(arr: &mut Array2<i32>, col_idx: usize) {
    assert!(col_idx < arr.ncols(), "Column index out of bounds");
    
    // 稳定降序排序
    arr.axis_iter_mut(Axis(0))
        .sort_by_key(|row| Reverse(row[col_idx]));
    
    // 不稳定降序排序（性能更优）：
    // arr.axis_iter_mut(Axis(0))
    //     .sort_unstable_by_key(|row| Reverse(row[col_idx]));
}

方法二：手动反转比较逻辑 #

通过sort_by直接反转列元素的比较顺序：

use ndarray::{Array2, Axis};
use std::cmp::Ordering;

fn sort_array_by_col_desc(arr: &mut Array2<i32>, col_idx: usize) {
    assert!(col_idx < arr.ncols(), "Column index out of bounds");
    
    arr.axis_iter_mut(Axis(0))
        .sort_by(|a, b| b[col_idx].cmp(&a[col_idx]));
}

4. 性能说明 #

• 稳定排序：sort_by_key/sort_by基于Timsort实现，会保留相等元素的原始顺序，适合需要稳定排序的场景。
• 不稳定排序：sort_unstable_by_key/sort_unstable_by采用快速排序变体，速度更快、内存占用更低，是性能优先场景的最优选择。

以上方法直接操作原数组内存，无需额外复制数据，是ndarray场景下的高效实现方式。

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