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

numpy.argsort https://docs.scipy.org/doc/numpy-1.13.0/reference/generated/numpy.argsort.html文档状态

Returns:
索引数组：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我想订购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

看起来这个功能是已经是 numpy 中的增强请求 https://github.com/numpy/numpy/issues/8708.

The numpy 问题 #8708 https://github.com/numpy/numpy/issues/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]])