Python中近乎无限维度二维数组的对角迭代优化方案问询

针对你需要从左上到右下按对角线顺序遍历数组、避免重复计算且可中途终止的需求，以下是高效的实现方案：

二维数组的最优遍历方式 #

核心逻辑是按**索引和i+j**分组遍历，每条分组对应一条从左上到右下的反对角线，每个元素仅被访问一次，无冗余计算。

代码实现 #

import numpy as np

# 生成你的测试数组
mat = np.zeros(shape=(10, 10), dtype='<U2')
for i in range(10):
    for j in range(10):
        mat[i][j] = str(i) + str(j)

n = mat.shape[0]  # 获取数组边长
# 遍历所有可能的索引和k（从0到2*(n-1)）
for k in range(0, 2 * (n - 1) + 1):
    # 确定当前k对应的i的合法范围
    start_i = max(0, k - (n - 1))
    end_i = min(n - 1, k)
    # 遍历当前反对角线上的所有元素
    for i in range(start_i, end_i + 1):
        j = k - i
        print((i, j), mat[i][j])
        
        # 加入阈值终止逻辑，示例：当元素为"33"时终止
        if mat[i][j] == "33":
            break
    # 若需要外层循环也终止，可设置flag
    # if 终止条件:
    #     break

方案优势 #

• 无重复计算：每个元素仅被访问一次，彻底解决了你之前方案中主对角线重复遍历的问题。
• 时间复杂度最优：O(n²)，与逐行遍历效率一致，但顺序完全符合需求。
• 终止逻辑灵活：可在任意元素处加入判断，满足阈值时直接终止迭代，避免冗余计算。

高维数组的扩展思路 #

如果是更高维度的数组（如d维），可按所有维度索引之和从小到大遍历，通过生成器生成合法索引组合：

import itertools

def diagonal_traverse(dim_size, num_dims):
    max_sum = num_dims * (dim_size - 1)
    for s in range(max_sum + 1):
        # 生成所有维度索引和为s的合法组合
        for indices in itertools.combinations_with_replacement(range(num_dims), s):
            cnt = [0] * num_dims
            for idx in indices:
                cnt[idx] += 1
            if all(c <= dim_size - 1 for c in cnt):
                yield tuple(cnt)

# 示例：遍历3x3x3的三维数组
for idx in diagonal_traverse(3, 3):
    print(idx)

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