运行时才可确定维度的n维数字数组遍历方法问询

坐标与线性索引映射原理 #

任意维度的数组，只要确定各维度长度组成的shape元组，即可通过步长数组实现线性索引和n维坐标的双向转换：

• 步长数组从最后一维开始计算，最后一维步长固定为1，前一维度的步长 = 后一维度步长 * 后一维度长度
• 示例：shape为(2,4,6,5)的数组，步长数组为[120, 30, 5, 1]

方案1：扁平化后线性索引转坐标 #

适合内存连续的存储场景，扁平化后遍历线性索引，按需转换为对应n维坐标：

def linear_to_coords(pos: int, shape: tuple) -> tuple:
    # 计算各维度步长
    strides = [1]
    for dim in reversed(shape[1:]):
        strides.insert(0, strides[0] * dim)
    # 拆解坐标
    coords = []
    remaining = pos
    for s in strides:
        coords.append(remaining // s)
        remaining %= s
    return tuple(coords)

调用示例：

shape = (2,4,6,5)
# 遍历所有元素
for pos in range(2*4*6*5):
    coords = linear_to_coords(pos, shape)
    # 处理业务逻辑，coords即为(i,j,k,l)格式的坐标

方案2：递归遍历直接获取坐标 #

不需要扁平化，递归过程携带前缀坐标，走到最后一层维度时即可拿到完整坐标和对应元素：

def traverse_ndarray(arr, prefix_coords: tuple = ()):
    # 判定是否到达最后一层数字元素
    if isinstance(arr, (int, float)):
        yield (prefix_coords, arr)
        return
    # 递归遍历当前维度所有索引
    for idx, sub_arr in enumerate(arr):
        yield from traverse_ndarray(sub_arr, prefix_coords + (idx,))

调用示例：

# 任意维度嵌套数字数组
test_arr = [[[[i*8 + j*4 + k*2 + l for l in range(2)] for k in range(2)] for j in range(2)] for i in range(2)]
for coords, val in traverse_ndarray(test_arr):
    print(f"坐标{coords}对应值：{val}")

方案3：迭代器实现（无递归深度限制） #

如果维度n极高，递归会触发栈溢出，用栈模拟递归过程即可规避：

def traverse_ndarray_iter(arr):
    stack = [(arr, ())]
    while stack:
        current, coords = stack.pop()
        if isinstance(current, (int, float)):
            yield (coords, current)
            continue
        # 倒序入栈保证遍历顺序与递归一致
        for idx in reversed(range(len(current))):
            stack.append((current[idx], coords + (idx,)))

工具库内置方案 #

如果使用Numpy数组，可直接调用内置接口实现，无需手动实现逻辑：

import numpy as np
# 任意维度numpy数组
arr = np.random.rand(2,4,6,5)
it = np.nditer(arr, flags=['multi_index'])
while not it.finished:
    coords = it.multi_index # 直接获取n维坐标
    val = it[0] # 获取对应元素值
    # 处理业务逻辑
    it.iternext()

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