用于多维 ndarray 的 argsort

Posted

技术标签:

【中文标题】用于多维 ndarray 的 argsort【英文标题】:argsort for a multidimensional ndarray 【发布时间】:2016-01-13 10:44:05 【问题描述】:

我正在尝试获取索引以按最后一个轴对多维数组进行排序,例如

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

【问题讨论】:

【参考方案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]])

【讨论】:

感谢@hpaulj,非常有用的解释!如果你有时间,你能解释一下'第一个维度的列索引[ing]'吗?那只是将数组转换为 (2,1,3) 对...为什么这有助于i 切片? 我扩展了我的解释。 有更简单的方法吗?我认为 argsort 应该在对数组进行排序后考虑它的用途?..... 除了在函数中隐藏细节之外,我不知道你希望什么样的简化。 argsort 可以以多种方式使用,而不仅仅是这种情况。将其视为一个构建块,您可以将其与其他一些块结合起来完成这项工作。 ***.com/questions/43220729/… 是另一个使用索引数组和arange 数组从行中选择元素的示例。同样的逻辑,只是索引数组不是由argsort生成的。【参考方案2】:

上面的答案现在有点过时了,因为在 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]])

【讨论】:

【参考方案3】:

我找到了the answer here,有人遇到了同样的问题。他们的关键只是欺骗索引以正常工作......

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

【讨论】:

哎呀我猜np.sort(dists, axis=1) 就是我要找的东西? 完全是@endolith。对于我的情况,我特别需要索引以相同的顺序对另一个数组进行排序。但我同意argsort 文档也可以使用更多改进;) 实际上我找到了一种无需任何排序的方法,只需沿轴使用 amin 和 amax【参考方案4】:

你也可以使用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

【讨论】:

以上是关于用于多维 ndarray 的 argsort的主要内容,如果未能解决你的问题,请参考以下文章

NumPy之:ndarray多维数组操作

Numpy:ndarray数据类型和运算

为啥使用数组作为索引会改变多维 ndarray 的形状?

Python——ndarray多维数组基本操作(1)

Numpy Ndarray 对象

Numpy Ndarray 对象