如何实现大型3D numpy数组的高效索引切片操作？

原代码lookup[:,m[:,1], m[:,0]]存在逻辑偏差：它会返回形状为(1000, 1000)的数组，相当于对每一张图像都提取了位置数组m里所有1000个坐标的值，总共读取100万条数据，这是耗时过高的核心原因。你实际需要的是对第i张图像仅提取m[i]对应的坐标值，仅需要读取1000条数据。

最优实现方案 #

方案1：对齐维度的花式索引（最直观，性能提升约300倍） #

给第0维也传入对应顺序的索引，让三个维度的索引一一匹配，直接拿到对应位置的1000个值：

import numpy as np
# 初始化测试数据
lookup = np.random.rand(1000,1000,1000)
m = (np.random.rand(1000,2)*1000).astype(int)

# 优化后写法
res = lookup[np.arange(len(lookup)), m[:,1], m[:,0]]

时间对比：

• 原写法：约20ms
• 优化后写法：约60μs，提升超过2个数量级

方案2：扁平化索引（适合超大数组场景，性能略高于方案1） #

先把三维坐标转成一维扁平化索引，直接从扁平化后的数组取值：

flat_idx = np.ravel_multi_index((np.arange(1000), m[:,1], m[:,0]), lookup.shape)
res = lookup.ravel()[flat_idx]

结果一致性验证 #

可以用np.allclose验证优化后结果和预期一致：

# 手动循环取的结果作为基准
benchmark = np.array([lookup[i, y, x] for i, (x,y) in enumerate(m)])
print(np.allclose(res, benchmark)) # 输出True

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