多维 ndarray 的 argsort
argsort for a multidimensional ndarray
我正在尝试获取索引以按最后一个轴对多维数组进行排序,例如
>>> 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]])
我发现 the answer here,有人遇到同样的问题。他们的关键只是欺骗索引正常工作...
>>> a[np.arange(np.shape(a)[0])[:,np.newaxis], np.argsort(a)]
array([[1, 2, 3],
[2, 8, 9]])
解决方案:
>>> 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])
要同时执行两个索引步骤,我们必须对第一个维度使用 'column' 索引。
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]])
作为 'short hand' 对于 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]])
你也可以使用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
上面的答案现在有点过时了,因为在 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]])
我正在尝试获取索引以按最后一个轴对多维数组进行排序,例如
>>> 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]])
我发现 the answer here,有人遇到同样的问题。他们的关键只是欺骗索引正常工作...
>>> a[np.arange(np.shape(a)[0])[:,np.newaxis], np.argsort(a)]
array([[1, 2, 3],
[2, 8, 9]])
解决方案:
>>> 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])
要同时执行两个索引步骤,我们必须对第一个维度使用 'column' 索引。
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]])
作为 'short hand' 对于 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]])
你也可以使用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
上面的答案现在有点过时了,因为在 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]])