具有多个维度的 numpy.argsort 的排序不变式

Question

numpy.argsort 文档状态

返回：
index_array ：ndarray，int 沿指定轴对 a 进行排序的索引数组。如果 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，我想使用np.argsort(A) 沿相应轴的结果订购B。考虑以下示例：

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

---

看起来这个功能是already an enhancement request in numpy。

Answer 1

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

Answer 2

这个 argsort 产生一个 (3,2) 数组

In [453]: idx=np.argsort(A,axis=-1)
In [454]: idx
Out[454]: 
array([[0, 1],
       [1, 0],
       [0, 1]], dtype=int32)

正如您所注意到的，将其应用于 A 以获得 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]])

虽然可能不是最快的，但它可能是最清晰的解决方案，也是构思更好解决方案的良好起点。

take 解决方案中的tuple(inds) 是：

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

第一部分是np.ix_ 将构造的，但它并不“喜欢”二维idx。

---

看起来我几年前探讨过这个话题

argsort for a multidimensional ndarray

a[np.arange(np.shape(a)[0])[:,np.newaxis], np.argsort(a)]

我试图解释发生了什么。 take 函数做同样的事情，但为更一般的情况（维度和轴）构造索引元组。推广到更多维度，但仍然使用axis=-1 应该很容易。

对于第一个轴，A[np.argsort(A,axis=0),np.arange(2)] 有效。

Answer 3

我们只需要使用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]])

Relevant one.