如何高效实现Numpy数组间的索引映射转换？

我有两个Numpy数组：一个3D或2D的reference数组，以及一个3D的map_数组。示例如下：

reference数组 #

import numpy as np
reference = np.array([[ [11,12,13], [14,15,16], [17,18,19] ],
                   [ [21,22,23], [24,25,26], [27,28,29] ],
                   [ [31,32,33], [34,35,36], [37,38,39] ],
                   [ [41,42,43], [44,45,46], [47,48,49] ]])

map_数组 #

map_ = np.array([[ [0,0], [1,2], [2,2] ],
              [ [2,1], [3,2], [3,0] ]])

期望结果 #

np.array([[ [11,12,13], [27,28,29], [37,38,39] ],
       [ [24,25,26], [47,48,49], [41,42,43] ]])

我用列表推导式实现了需求：

np.array([[reference[i[0],i[1],:] for i in j] for j in map_])

但当数组规模达到数千级别时，这个方法运行速度极慢，有没有更高效的实现方式？

使用Numpy的向量化高级索引可以彻底解决列表推导式的性能问题，所有操作都在Numpy内部的C层执行，避免Python循环的开销。

实现步骤 #

1. 从map_数组中拆分出对应的行索引和列索引：
   • rows = map_[..., 0]：提取每个映射对的第一个元素作为行索引
   • cols = map_[..., 1]：提取每个映射对的第二个元素作为列索引
2. 直接用这两个索引数组对reference进行索引，Numpy会自动广播匹配形状，生成目标数组。

完整代码示例 #

import numpy as np

reference = np.array([[ [11,12,13], [14,15,16], [17,18,19] ],
                   [ [21,22,23], [24,25,26], [27,28,29] ],
                   [ [31,32,33], [34,35,36], [37,38,39] ],
                   [ [41,42,43], [44,45,46], [47,48,49] ]])

map_ = np.array([[ [0,0], [1,2], [2,2] ],
              [ [2,1], [3,2], [3,0] ]])

# 拆分索引
rows = map_[..., 0]
cols = map_[..., 1]

# 向量化索引
result = reference[rows, cols]

print(result)

输出结果 #

[[[11 12 13]
  [27 28 29]
  [37 38 39]]

 [[24 25 26]
  [47 48 49]
  [41 42 43]]]

适用场景 #

• 无论reference是2D（元素为标量）还是3D（元素为向量），该方法都能直接生效
• 对于规模达数千甚至更大的数组，性能比列表推导式提升10~100倍以上

内容的提问来源于stack exchange，提问作者Carlos Eugenio Thompson Pinzón