Rust ndarray中如何按负步长迭代原数组索引？

ndarray当前没有直接提供API，让你基于SliceInfo直接迭代原数组的对应索引（而非切片后的本地索引），但可以通过解析SliceInfo的维度参数手动实现需求。

实现思路 #

1. 解析SliceInfo中每个维度的切片规则，调用slice.indices(dim_size)获取该维度在原数组中的(start, stop, step)参数。
2. 根据步长正负，生成对应维度的原索引序列：
   • 正步长：直接用start..stop按步长生成序列
   • 负步长：从start递减到stop+1（因为range是左闭右开），按步长绝对值生成反向序列
3. 将各维度的索引序列做笛卡尔积，得到符合要求的原数组索引迭代器。

代码示例 #

use ndarray::{Array2, SliceInfo, s, Ix2};
use itertools::Itertools;

fn main() {
    // 示例2D数组
    let arr = Array2::zeros((20, 20));
    // 目标SliceInfo
    let slice_spec = s![5..10, 10..20;-1];
    let slice_info = SliceInfo::new(slice_spec).unwrap();

    // 生成每个维度的原索引序列
    let dim_index_sequences: Vec<_> = slice_info
        .iter()
        .enumerate()
        .map(|(dim_idx, slice)| {
            let dim_size = arr.raw_dim()[dim_idx];
            let (start, stop, step) = slice.indices(dim_size);
            
            if step > 0 {
                (start..stop).step_by(step).collect::<Vec<_>>()
            } else {
                // 负步长时，范围是从start到stop+1（左闭右开），反向遍历
                (stop + 1..=start).rev().step_by(step.abs()).collect::<Vec<_>>()
            }
        })
        .collect();

    // 生成原数组索引迭代器
    let original_indices_iter = dim_index_sequences[0]
        .iter()
        .cartesian_product(dim_index_sequences[1].iter())
        .map(|(i, j)| (*i, *j));

    // 验证输出：会依次输出(5,19), (5,18)...(5,10), (6,19)...(6,10)...
    for idx in original_indices_iter {
        println!("{:?}", idx);
    }
}

说明 #

• 这个方法支持任意维度的SliceInfo，只要替换数组和切片规则即可适配。
• 如果不想引入itertools，可以手动嵌套循环来生成笛卡尔积，比如外层遍历第一个维度的索引，内层遍历第二个维度的反向索引并组合。

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