numpy数组切片，从每个第三维获取一个

Question

我有一个 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]])

有没有简洁+快速的方法？

Answer 1

如果我理解正确的话，我已经使用花哨的索引做了类似的事情：

>>> k,j = np.meshgrid(np.arange(a.shape[1]),np.arange(a.shape[0]))
>>> k
array([[0, 1, 2, 3, 4],
       [0, 1, 2, 3, 4],
       [0, 1, 2, 3, 4]])
>>> j
array([[0, 0, 0, 0, 0],
       [1, 1, 1, 1, 1],
       [2, 2, 2, 2, 2]])
>>> a[j,k,getters]
array([[  0,   8,  16,  21,  29],
       [ 37,  42,  50,  58,  63],
       [ 71,  79,  84,  92, 100]])

当然，您可以保留k 和j 并随时使用它们。正如 DSM 在下面的 cmets 中指出的那样，j,k = np.indices(a.shape[:2]) 也应该代替meshgrid 起作用。哪个更快（显然）取决于您使用的元素数量。