如何使用3分量索引列表从3D NumPy数组中提取对应值

对于你的需求，Numpy的高级索引可以直接实现简洁高效的解决方案，完全不需要循环或手动转换索引维度：

import numpy as np
data = np.arange(125).reshape(5,5,5)
indices = np.array(((0, 0, 0), (4, 4, 4), (0, 0, 4)), dtype=int)

# 核心简洁写法
values = data[indices[:, 0], indices[:, 1], indices[:, 2]]
print(values)  # 输出: [  0 124   4]

原理说明 #

Numpy支持将多个一维数组作为索引参数，分别对应目标数组的各个维度。这里：

• indices[:, 0]提取所有索引的x分量（第一维度索引）
• indices[:, 1]提取所有索引的y分量（第二维度索引）
• indices[:, 2]提取所有索引的z分量（第三维度索引）

将这三个一维数组传入data的索引位置，就能一次性匹配并提取所有对应位置的值，效率远高于循环，写法也比扁平化数组的方式简洁得多。

更直观的写法 #

如果觉得拆分索引的写法不够直观，还可以通过转置索引数组后解包，逻辑和手动写data[ix, iy, iz]完全一致：

values = data[*np.transpose(indices)]

对比原有方案 #

• 循环写法：逻辑直观但大数据量下效率极低，Numpy不推荐用循环处理数组
• 扁平化数组写法：需要额外编写索引转换函数，步骤冗余且可读性差
• 高级索引写法：Numpy原生支持，代码简洁、效率高，是这类场景的最优解

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