如何无需循环访问3D NumPy数组中的指定多列？

当然可以！你现在用循环的方式虽然可行，但对于大型数组来说效率不够高，NumPy的高级索引正好能帮你摆脱循环，而且代码更简洁。

解决方案代码 #

直接利用数组索引的广播特性，一步就能提取所有需要的列：

import numpy as np

# 假设已经定义好数组A和B
i_indices = B[:, 0]
j_indices = B[:, 1]
# 一次性提取所有目标列，结果形状为(N, K)，K是B的行数
data_out_array = A[:, i_indices, j_indices]
# 如果需要和原循环输出的(K, N)结构一致，转置即可
data_out = data_out_array.T
# 若确实需要列表形式（不推荐，保持数组更高效），可转换为列表
data_out_list = data_out_array.T.tolist()

为什么这能行？ #

原循环里每次取A[:, i, j]，本质是在第1维度（M维）取i、第2维度（L维）取j。用高级索引时，i_indices和j_indices都是长度为K的一维数组，NumPy会自动将它们配对，分别在第1、第2维度上对应索引，一次性提取所有A[:, i_k, j_k]。最终得到的A[:, i_indices, j_indices]是形状为(N, K)的数组，转置后就和你循环append得到的(K, N)结构完全一致。

验证结果一致性 #

可以用下面的代码确认优化后的结果和原循环完全相同：

# 原循环代码
data_out_loop = []
for p in B:
    i, j = p
    col = A[:, i, j]
    data_out_loop.append(col)
data_out_loop = np.array(data_out_loop)

# 对比结果
assert np.allclose(data_out, data_out_loop)

额外提示 #

• 确保B里的索引都是整数类型，避免索引时出现类型错误；
• 高级索引是NumPy底层优化的操作，比Python循环快得多，尤其是当B的行数很多时，性能提升会非常明显；
• 尽量保持结果为NumPy数组而不是列表，后续的数值计算效率会更高。

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