我还没有想到如何在N维中进行此操作,但这
是2D版本:
>>> a = np.random.standard_normal(size=(2,5))
>>> a
array([[ 0.72322499, -0.05376714, -0.28316358, 1.43025844, -0.90814293],
[ 0.7459107 , 0.43020728, 0.05411805, -0.32813465, 2.38829386]])
>>> i = np.array([[0,1,2,4,3],[0,1,2,3,4]])
>>> a[np.arange(a.shape[0])[:,np.newaxis],i]
array([[ 0.72322499, -0.05376714, -0.28316358, -0.90814293, 1.43025844],
[ 0.7459107 , 0.43020728, 0.05411805, -0.32813465, 2.38829386]])
这是N维版本:
>>> a[list(np.ogrid[[slice(x) for x in a.shape]][:-1])+[i]]
运作方式如下:
好的,让我们从3维数组开始进行说明。
>>> import numpy as np
>>> a = np.arange(24).reshape((2,3,4))
>>> a
array([[[ 0, 1, 2, 3],
[ 4, 5, 6, 7],
[ 8, 9, 10, 11]],
[[12, 13, 14, 15],
[16, 17, 18, 19],
[20, 21, 22, 23]]])
您可以通过沿每个轴指定索引来访问此数组的元素,如下所示:
>>> a[0,1,2]
6
这等效于a[0][1][2]
如果我们要处理列表而不是数组,您将如何访问同一元素。
切片数组时,Numpy可以使您变得更加幻想:
>>> a[[0,1],[1,1],[2,2]]
array([ 6, 18])
>>> a[[0,1],[1,2],[2,2]]
array([ 6, 22])
这些例子就等于[a[0][1][2],a[1][1][2]]
和[a[0][1][2],a[1][2][2]]
如果我们处理列表。
您甚至可以省去重复的索引,而numpy会找出您想要的内容。例如,以上示例可以等效地编写:
>>> a[[0,1],1,2]
array([ 6, 18])
>>> a[[0,1],[1,2],2]
array([ 6, 22])
在每个维度中切片的数组(或列表)的 形状 仅影响返回数组的 形状 。换句话说,numpy不在乎您是否在尝试(2,3,4)
拉取值时尝试使用形状数组来索引数组,只是它会反馈形状数组(2,3,4)
。例如:
>>> a[[[0,0],[0,0]],[[0,0],[0,0]],[[0,0],[0,0]]]
array([[0, 0],
[0, 0]])
在这种情况下,我们a[0,0,0]
一遍又一遍地抓取相同的元素,但是numpy返回的数组与传入的形状相同。
好的,解决您的问题。您想要的是沿最后一个轴用index
数组中的数字索引数组。因此,对于您问题中的示例,您希望[[a[0,0],a[0,1],a[0,2],a[0,4],a[0,3]],a[1,0],a[1,1],...
正如我前面所说,索引数组是多维的,这一事实并不能告诉numpy有关从何处提取这些索引的任何信息。它只是指定输出数组的形状。因此,在您的示例中,您需要告诉numpy,要从中提取前5个值,而从中提取a[0]
后5个值a[1]
。简单!
>>> a[[[0]*5,[1]*5],index]
它在N维中变得复杂,但让我们对a
我上面定义的3维数组进行处理。假设我们有以下索引数组:
>>> i = np.array(range(4)[::-1]*6).reshape(a.shape)
>>> i
array([[[3, 2, 1, 0],
[3, 2, 1, 0],
[3, 2, 1, 0]],
[[3, 2, 1, 0],
[3, 2, 1, 0],
[3, 2, 1, 0]]])
因此,这些值全部用于沿最后一个轴的索引。我们需要告诉numpy这些数字将沿第一和第二轴取哪些索引;即我们需要告诉numpy第一个轴的索引是:
i1 = [[[0, 0, 0, 0],
[0, 0, 0, 0],
[0, 0, 0, 0]],
[[1, 1, 1, 1],
[1, 1, 1, 1],
[1, 1, 1, 1]]]
第二轴的索引为:
i2 = [[[0, 0, 0, 0],
[1, 1, 1, 1],
[2, 2, 2, 2]],
[[0, 0, 0, 0],
[1, 1, 1, 1],
[2, 2, 2, 2]]]
然后我们可以这样做:
>>> a[i1,i2,i]
array([[[ 3, 2, 1, 0],
[ 7, 6, 5, 4],
[11, 10, 9, 8]],
[[15, 14, 13, 12],
[19, 18, 17, 16],
[23, 22, 21, 20]]])
方便的numpy函数,它生成i1
并i2
称为np.mgrid
。我np.ogrid
在答案中使用了这种方法,在这种情况下,这是等效的,因为我之前谈到过麻木的魔术。
希望有帮助!