np.nditer用法理解及遍历切片时元素原索引获取问题

待分析代码 #

import numpy as np
arr = np.array([[1, 2, 3, 4,99,11,22], [5, 6, 7, 8,43,54,22]])
for x in np.nditer(arr[0:,::4]):
  print(x)

对你现有理解的校验 #

• 你提到的「二维数组包含两个一维子数组」是正确的，该数组形状为(2,7)，即2行7列。
• 你对np.nditer(arr[0:,::4])的理解不完全准确：
  • 逗号分隔的切片规则分别对应NumPy数组的两个维度，第一个维度（行方向）的0:确实表示取从索引0到末尾的所有行，也就是该数组的全部2行。
  • 第二个维度（列方向）的::4不是遍历打印的步长，而是切片阶段就生效的列筛选规则：切片语法[起始:结束:步长]省略起始、结束值时默认从维度头取到维度尾，步长为4，即选出列索引满足「从0开始每4个取1个」的列。
  • 最终切出的子数组形状为(2,2)，内容是[[1,99],[5,43]]，np.nditer只是按内存顺序遍历这个已经切好的子数组的所有元素，实际打印输出为1、99、5、43。

带步长切片场景下获取原数组索引的方法 #

直接对切片生成的子数组用np.nditer遍历，只能拿到子数组的元素值，不会自带原数组的索引信息，你可以用以下两种方法获取精确原索引：

方法1：基于已知切片规则手动计算偏移 #

切片规则明确时，遍历子数组的内部索引，按切片的起始位置、步长反推原数组索引即可，效率最高：

# 定义切片规则，和你写在[]里的规则保持一致
row_start = 0
col_start = 0
col_step = 4
sliced_arr = arr[row_start::, col_start::col_step]

# 用ndenumerate同时获取子数组内部索引和对应值
for (row_in_slice, col_in_slice), val in np.ndenumerate(sliced_arr):
    origin_row = row_start + row_in_slice
    origin_col = col_start + col_in_slice * col_step
    print(f"当前值：{val}，原数组索引：({origin_row}, {origin_col})")

运行输出：

当前值：1，原数组索引：(0, 0)
当前值：99，原数组索引：(0, 4)
当前值：5，原数组索引：(1, 0)
当前值：43，原数组索引：(1, 4)

方法2：索引掩码筛选 #

如果切片规则比较复杂，不想手动计算偏移，可以直接生成原数组的全量索引网格，按切片规则生成掩码筛选出符合条件的索引和值：

# 生成原数组所有位置的行、列索引网格
rows, cols = np.indices(arr.shape)
# 按切片规则写筛选条件：行>=0（所有行）、列索引从0开始步长为4
mask = (rows >= 0) & ((cols - 0) % 4 == 0)
# 遍历所有符合条件的位置
for r, c, val in zip(rows[mask], cols[mask], arr[mask]):
    print(f"当前值：{val}，原数组索引：({r}, {c})")

该方法运行结果和方法1完全一致，不需要手动换算偏移，适配复杂切片场景。

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