Question

我试图让索引按最后一个轴对多维数组进行排序，例如

>>> a = np.array([[3,1,2],[8,9,2]])

我喜欢索引i，

>>> a[i]
array([[1, 2, 3],
       [2, 8, 9]])

根据numpy.argsort的文档，我认为应该这样做，但我收到了错误：

>>> a[np.argsort(a)]
IndexError: index 2 is out of bounds for axis 0 with size 2

编辑：我需要以相同的方式重新排列相同形状的其他数组（例如数组b，使a.shape == b.shape）......以便

>>> b = np.array([[0,5,4],[3,9,1]])
>>> b[i]
array([[5,4,0],
       [9,3,1]])

Answer 1

解决方案：

>>> a[np.arange(np.shape(a)[0])[:,np.newaxis], np.argsort(a)]
array([[1, 2, 3],
       [2, 8, 9]])

你做对了，虽然我不会把它描述为欺骗索引。

也许这会有助于使其更清晰：

In [544]: i=np.argsort(a,axis=1)

In [545]: i
Out[545]: 
array([[1, 2, 0],
       [2, 0, 1]])

对于每一行，

i是我们想要的顺序。那就是：

In [546]: a[0, i[0,:]]
Out[546]: array([1, 2, 3])

In [547]: a[1, i[1,:]]
Out[547]: array([2, 8, 9])

要同时执行两个索引步骤，我们必须为第一维使用“列”索引。

In [548]: a[[[0],[1]],i]
Out[548]: 
array([[1, 2, 3],
       [2, 8, 9]])

可以与i配对的另一个数组是：

In [560]: j=np.array([[0,0,0],[1,1,1]])

In [561]: j
Out[561]: 
array([[0, 0, 0],
       [1, 1, 1]])

In [562]: a[j,i]
Out[562]: 
array([[1, 2, 3],
       [2, 8, 9]])

如果i标识每个元素的列，则j指定每个元素的行。 [[0],[1]]列数组也可以正常工作，因为它可以针对i进行广播。

我想到了

np.array([[0],
          [1]])

作为j的'简称'。它们一起定义新数组的每个元素的源行和列。他们一起工作，而不是按顺序工作。

从a到新数组的完整映射是：

[a[0,1]  a[0,2]  a[0,0]
 a[1,2]  a[1,0]  a[1,1]]

def foo(a):
    i = np.argsort(a, axis=1)
    return (np.arange(a.shape[0])[:,None], i)

In [61]: foo(a)
Out[61]: 
(array([[0],
        [1]]), array([[1, 2, 0],
        [2, 0, 1]], dtype=int32))
In [62]: a[foo(a)]
Out[62]: 
array([[1, 2, 3],
       [2, 8, 9]])

Answer 2

我找到{{3}}，有人遇到同样的问题。他们的关键是欺骗索引以正常工作......

repaint

Answer 3

你也可以使用linear indexing，这可能会提高性能，就像这样 -

M,N = a.shape
out = b.ravel()[a.argsort(1)+(np.arange(M)[:,None]*N)]

因此，a.argsort(1)+(np.arange(M)[:,None]*N)基本上是用于映射b以获得b所需排序输出的线性索引。也可以在a上使用相同的线性索引来获取a的排序输出。

示例运行 -

In [23]: a = np.array([[3,1,2],[8,9,2]])

In [24]: b = np.array([[0,5,4],[3,9,1]])

In [25]: M,N = a.shape

In [26]: b.ravel()[a.argsort(1)+(np.arange(M)[:,None]*N)]
Out[26]: 
array([[5, 4, 0],
       [1, 3, 9]])

Rumtime测试 -

In [27]: a = np.random.rand(1000,1000)

In [28]: b = np.random.rand(1000,1000)

In [29]: M,N = a.shape

In [30]: %timeit b[np.arange(np.shape(a)[0])[:,np.newaxis], np.argsort(a)]
10 loops, best of 3: 133 ms per loop

In [31]: %timeit b.ravel()[a.argsort(1)+(np.arange(M)[:,None]*N)]
10 loops, best of 3: 96.7 ms per loop

Answer 4

由于numpy 1.15中添加了新功能以使其更简单，因此上述答案现在已过时。 take_along_axis（https://docs.scipy.org/doc/numpy-1.15.1/reference/generated/numpy.take_along_axis.html）允许您执行以下操作：

>>> a = np.array([[3,1,2],[8,9,2]])
>>> np.take_along_axis(a, a.argsort(axis=-1), axis=-1)
array([[1 2 3]
       [2 8 9]])

argsort为多维ndarray

4 个答案: