对具有多个维度的numpy.argsort进行不变排序

对具有多个维度的numpy.argsort进行不变排序

该numpy的问题＃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]])