我有一个 3D 数据数组。我有一个二维索引数组,其中形状与数据数组的前两个维度相匹配,并且它指定了我想从数据数组中提取的索引以创建一个二维数组。例如:
from numpy import *
a = arange(3 * 5 * 7).reshape((3,5,7))
getters = array([0,1,2] * (5)).reshape(3,5)
我正在寻找的是一种语法a[:, :, getters]
,它通过独立索引到每个项目的第三维来返回一个形状为 (3,5) 的数组。但是,a[:, :, getters]
返回一个形状为 (3,5,3,5) 的数组。我可以通过迭代和构建一个新数组来做到这一点,但这很慢:
array([[col[getters[ri,ci]] for ci,col in enumerate(row)] for ri,row in enumerate(a)])
# gives array([[ 0, 8, 16, 21, 29],
# [ 37, 42, 50, 58, 63],
# [ 71, 79, 84, 92, 100]])
有没有简洁+快速的方法?