numpy.argsort
docs state
返回:
index_array:ndarray,int 沿指定轴排序的索引数组。如果a是一维的,a[index_array]
会产生一个排序的a。
如何将numpy.argsort
的结果应用于多维数组以获取已排序的数组? (不只是一维或二维数组;它可能是一个N维数组,其中N只在运行时知道)
>>> import numpy as np
>>> np.random.seed(123)
>>> A = np.random.randn(3,2)
>>> A
array([[-1.0856306 , 0.99734545],
[ 0.2829785 , -1.50629471],
[-0.57860025, 1.65143654]])
>>> i=np.argsort(A,axis=-1)
>>> A[i]
array([[[-1.0856306 , 0.99734545],
[ 0.2829785 , -1.50629471]],
[[ 0.2829785 , -1.50629471],
[-1.0856306 , 0.99734545]],
[[-1.0856306 , 0.99734545],
[ 0.2829785 , -1.50629471]]])
对我而言,这不仅仅是使用sort()
的问题;我有另一个数组B
,我想在适当的轴上使用B
的结果订购np.argsort(A)
。请考虑以下示例:
>>> A = np.array([[3,2,1],[4,0,6]])
>>> B = np.array([[3,1,4],[1,5,9]])
>>> i = np.argsort(A,axis=-1)
>>> BsortA = ???
# should result in [[4,1,3],[5,1,9]]
# so that corresponding elements of B and sort(A) stay together
答案 0 :(得分:3)
numpy issue #8708有一个take_along_axis的示例实现,可以满足我的需要;我不确定它对大型阵列是否有效,但似乎有用。
def take_along_axis(arr, ind, axis):
"""
... here means a "pack" of dimensions, possibly empty
arr: array_like of shape (A..., M, B...)
source array
ind: array_like of shape (A..., K..., B...)
indices to take along each 1d slice of `arr`
axis: int
index of the axis with dimension M
out: array_like of shape (A..., K..., B...)
out[a..., k..., b...] = arr[a..., inds[a..., k..., b...], b...]
"""
if axis < 0:
if axis >= -arr.ndim:
axis += arr.ndim
else:
raise IndexError('axis out of range')
ind_shape = (1,) * ind.ndim
ins_ndim = ind.ndim - (arr.ndim - 1) #inserted dimensions
dest_dims = list(range(axis)) + [None] + list(range(axis+ins_ndim, ind.ndim))
# could also call np.ix_ here with some dummy arguments, then throw those results away
inds = []
for dim, n in zip(dest_dims, arr.shape):
if dim is None:
inds.append(ind)
else:
ind_shape_dim = ind_shape[:dim] + (-1,) + ind_shape[dim+1:]
inds.append(np.arange(n).reshape(ind_shape_dim))
return arr[tuple(inds)]
产生
>>> A = np.array([[3,2,1],[4,0,6]])
>>> B = np.array([[3,1,4],[1,5,9]])
>>> i = A.argsort(axis=-1)
>>> take_along_axis(A,i,axis=-1)
array([[1, 2, 3],
[0, 4, 6]])
>>> take_along_axis(B,i,axis=-1)
array([[4, 1, 3],
[5, 1, 9]])
答案 1 :(得分:1)
这个argsort产生一个(3,2)数组
map.locate
正如您所说,将此应用于In [453]: idx=np.argsort(A,axis=-1)
In [454]: idx
Out[454]:
array([[0, 1],
[1, 0],
[0, 1]], dtype=int32)
以获得相当于A
的信息并不明显。迭代解决方案是每行(1d情况)排序:
np.sort(A, axis=-1)
尽管可能不是最快的,但它可能是最清晰的解决方案,也是构建更好解决方案的良好起点。
In [459]: np.array([x[i] for i,x in zip(idx,A)])
Out[459]:
array([[-1.0856306 , 0.99734545],
[-1.50629471, 0.2829785 ],
[-0.57860025, 1.65143654]])
解决方案中的tuple(inds)
是:
take
换句话说:
(array([[0],
[1],
[2]]),
array([[0, 1],
[1, 0],
[0, 1]], dtype=int32))
In [470]: A[_]
Out[470]:
array([[-1.0856306 , 0.99734545],
[-1.50629471, 0.2829785 ],
[-0.57860025, 1.65143654]])
第一部分是In [472]: A[np.arange(3)[:,None], idx]
Out[472]:
array([[-1.0856306 , 0.99734545],
[-1.50629471, 0.2829785 ],
[-0.57860025, 1.65143654]])
将构建的内容,但它并不像&#39;喜欢&#39; 2d np.ix_
。
看起来我几年前探讨过这个话题
argsort for a multidimensional ndarray
idx
我试图解释发生了什么。 a[np.arange(np.shape(a)[0])[:,np.newaxis], np.argsort(a)]
函数执行相同的操作,但为更一般的情况(维度和轴)构造索引元组。推广到更多维度,但仍然使用take
应该很容易。
对于第一个轴,axis=-1
有效。
答案 2 :(得分:1)
我们只需要使用advanced-indexing
沿着所有轴索引那些索引数组。我们可以使用np.ogrid
沿所有轴创建范围数组的开放网格,然后仅使用输入索引替换输入轴。最后,索引到具有所需输出的索引的数据数组。因此,基本上,我们会 -
# Inputs : arr, ind, axis
idx = np.ogrid[tuple(map(slice, ind.shape))]
idx[axis] = ind
out = arr[tuple(idx)]
为了使其正常运行并进行错误检查,让我们创建两个函数 - 一个用于获取这些索引,另一个用于输入数据数组并简单地索引。第一个函数的想法是获取可以重新索引的索引,以便索引到任意数组,这些数组将支持沿每个轴的必要数量的维度和长度。
因此,实现将是 -
def advindex_allaxes(ind, axis):
axis = np.core.multiarray.normalize_axis_index(axis,ind.ndim)
idx = np.ogrid[tuple(map(slice, ind.shape))]
idx[axis] = ind
return tuple(idx)
def take_along_axis(arr, ind, axis):
return arr[advindex_allaxes(ind, axis)]
样品运行 -
In [161]: A = np.array([[3,2,1],[4,0,6]])
In [162]: B = np.array([[3,1,4],[1,5,9]])
In [163]: i = A.argsort(axis=-1)
In [164]: take_along_axis(A,i,axis=-1)
Out[164]:
array([[1, 2, 3],
[0, 4, 6]])
In [165]: take_along_axis(B,i,axis=-1)
Out[165]:
array([[4, 1, 3],
[5, 1, 9]])