Python导入三维数组Pe时自动循环遍历Pe.Pe[i]索引的实现方法

你持有维度为(2,3,3)的数组Pe，导入该数组后需要对每个沿第0轴切分的子数组自动执行遍历计算，返回对应的已访问索引。当前手动指定Pe.Pe[0]、Pe.Pe[1]传参的方式灵活性差，需要实现i在合法取值范围内自动传入Pe.Pe[i]完成遍历。

Pe数组定义 #

import numpy as np
Pe = np.array([[[128.22918457, 168.52413295, 209.72343319],
        [129.01598287, 179.03716051, 150.68633749],
        [131.00688309, 187.42601593, 193.68172751]],

       [[ 87.70103267, 115.2603484 , 143.4381863 ],
        [ 88.23915528, 122.45062554, 103.06037156],
        [ 89.60081102, 128.18809696, 132.46662659]]])
print([Pe])

功能实现代码 #

import numpy as np
import time
import Pe

def get_neighbor_indices(position, dimensions):
    '''
    dimensions is a shape of np.array
    '''
    i, j = position
    indices = [(i+1,j), (i-1,j), (i,j+1), (i,j-1)]
    return [
        (i,j) for i,j in indices
        if i>=0 and i<dimensions[0]
            and j>=0 and j<dimensions[1]
        ]

def iterate_array(init_i, init_j, arr, condition_func):
    '''
    arr is an instance of np.array
    condition_func is a function (value) => boolean
    '''
    indices_to_check = [(init_i,init_j)]
    checked_indices = set()
    result = []
    t0 = None
    t1 = None
    timestamps = []
    while indices_to_check:
        pos = indices_to_check.pop()
        if pos in checked_indices:
            continue
        item = arr[pos]
        checked_indices.add(pos)
        if condition_func(item):
            result.append(item)
            t1=time.time()
            if(t0==None):
                t0=t1
            timestamps.append(t1-t0)
            indices_to_check.extend(
                get_neighbor_indices(pos, arr.shape)
            )
    return result,timestamps


Visited_Indices,timestamps=iterate_array(0,0, Pe.Pe[0], lambda x : x < 150)
out = list(zip(*np.where(np.isin(Pe, Visited_Indices))))
print("Visited =",[Visited_Indices]) 

直接通过for循环遍历第0维的所有合法索引即可，不需要手动硬编码0、1两个索引值。用Pe.Pe.shape[0]获取数组第0维的长度，后续如果第0维的子数组数量变化，代码不需要修改就能自动适配。

将代码末尾的手动调用部分替换为如下逻辑：

# 初始化列表存储所有子数组的计算结果
all_visited_indices = []
all_timestamps = []
all_out = []

# 自动遍历第0维所有子数组
for idx in range(Pe.Pe.shape[0]):
    current_sub_arr = Pe.Pe[idx]
    visited, ts = iterate_array(0, 0, current_sub_arr, lambda x: x < 150)
    out = list(zip(*np.where(np.isin(current_sub_arr, visited))))
    
    # 存入结果集
    all_visited_indices.append(visited)
    all_timestamps.append(ts)
    all_out.append(out)
    
    # 打印当前子数组结果
    print(f"第{idx}个子数组Visited =", visited)

注意：原代码中np.isin(Pe, Visited_Indices)是在整个三维数组上做匹配，会匹配到不同子数组里相同的数值，建议改为在当前遍历的子数组上做匹配，避免结果混淆。

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