Numpy:通过多维数组对多维数组进行排序

Question

如果这是多余的或超级基本的,请原谅我.我是从R来到Python/Numpy并且很难在脑子里翻看东西.

我有一个维度数组,我想使用另一个索引值的n维数组进行排序.我知道我可以将它包装在一个循环中,但似乎应该有一个非常简洁的Numpyonic方式来击败这个提交.这是我设置n = 2的问题的示例代码:

a1 = random.standard_normal(size=[2,5]) 
index = array([[0,1,2,4,3] , [0,1,2,3,4] ]) 

所以现在我有一个2 x 5的随机数组和一个2 x 5的索引.我现在已经阅读了take()大约10次的帮助,但很明显,我的大脑并没有去研究它.

我以为这可能会让我在那里:

take(a1, index)

array([[ 0.29589188, -0.71279375, -0.18154864, -1.12184984,  0.25698875],
       [ 0.29589188, -0.71279375, -0.18154864,  0.25698875, -1.12184984]])

但这显然只重新排序了第一个元素(我认为是因为扁平化).

有关我如何从我所处的位置获得的解决方案,该解决方案将索引的元素0对a1的元素0进行排序...元素n？

Answer 1

我想不出如何在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]]

以下是它的工作原理:

好的,让我们从三维数组开始进行说明.

>>> 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要从中拉出前5个值,从中拉出a[0]后5个a[1].简单!

>>> a[[[0]*5,[1]*5],index]

它在N维中变得复杂,但让我们为a我上面定义的三维数组做这件事.假设我们有以下索引数组:

>>> 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]]])

生成i1和i2调用的方便numpy函数np.mgrid.我np.ogrid在我的答案中使用了相当于这种情况的答案,因为我之前谈过的那种笨拙的魔法.

希望有所帮助!