我有一个维度为 n 的数组a，以及一个维度为 n-1 的数组b。 b 最后一个轴的值对应于我要从数组 a 中提取的值的索引，在维度为 n-1 的数组res中。
对于 n=2 的示例：
a = np.array([[1, 2, 3, 4],
[5, 6, 7, 8],
[9, 10, 11, 12])
b = np.array([1,3,0])
我想要
res = [a[1], a[3], a[0]]
# i.e. res = [2, 8, 9]
有没有一个功能可以有效地做到这一点，具有更高的维度数？我知道我可以使用for循环，但我希望有更有效的循环。
编辑：
当 n=3 时，设a的形状为(2,2,3)。
然后b和res有(2,2)的形状：
a = np.array([[[ 1, 2, 3],
[ 4, 5, 6]],
[[ 7, 8, 9],
[10,11,12]]])
b = np.array([[0,2],
[1,2]]
# res = np.array([[1,6],
# [8,12]])
最新的numpy(1.15(增加了一个take_along_axis函数：
In [36]: np.take_along_axis(a, b[:,None], 1)
Out[36]:
array([[2],
[8],
[9]])
它使用帮助程序函数来构造索引元组：
In [37]: np.lib.shape_base._make_along_axis_idx((3,4), b[:,None], 1)
Out[37]:
(array([[0],
[1],
[2]]),
array([[1],
[3],
[0]]))
在此之前，我(和其他人(会建议：
In [38]: a[np.arange(3), b]
Out[38]: array([2, 8, 9])
这本质上是一回事(除了增加的维度(。 正如take_along_axis文档所示，这是为了沿轴获取argsort的结果之类的东西而设计的。
对于更高维度的情况：
In [39]: a1 = np.array([[[ 1, 2, 3],
...: [ 4, 5, 6]],
...: [[ 7, 8, 9],
...: [10,11,12]]])
...: b1 = np.array([[0,2],
...: [1,2]])
In [40]: a1.shape
Out[40]: (2, 2, 3)
In [41]: b1.shape
Out[41]: (2, 2)
In [42]: np.take_along_axis(a1, b1[...,None], -1)
Out[42]:
array([[[ 1],
[ 6]],
[[ 8],
[12]]])
In [45]: np.lib.shape_base._make_along_axis_idx(a1.shape, b1[...,None], 2)
Out[45]:
(array([[[0]],
[[1]]]),
array([[[0],
[1]]]),
array([[[0],
[2]],
[[1],
[2]]]))
In [46]: [i.shape for i in _]
Out[46]: [(2, 1, 1), (1, 2, 1), (2, 2, 1)]
同样，等效的自己动手索引：
In [48]: a1[np.arange(2)[:,None], np.arange(2)[None,:], b1]
Out[48]:
array([[ 1, 6],
[ 8, 12]])
一旦您了解了阵列广播及其如何应用于索引，这里的概念就不难了。 但是take_along_axis可能会使应用它们更容易。从某种意义上说，它是np.ix_的延伸。
In [50]: np.ix_(np.arange(2), np.arange(2), np.arange(3))
Out[50]:
(array([[[0]],
[[1]]]), array([[[0],
[1]]]), array([[[0, 1, 2]]]))
In [51]: [i.shape for i in _]
Out[51]: [(2, 1, 1), (1, 2, 1), (1, 1, 3)]
In [55]: a1[(*np.ix_(np.arange(2), np.arange(2)),b1)]
Out[55]:
array([[ 1, 6],
[ 8, 12]])