How to use numpy.argsort() as indices in more than 2 dimensions?

问题:

I know something similar to this question has been asked many times over already, but all answers given to similar questions only seem to work for arrays with 2 dimensions.

My understanding of np.argsort() is that np.sort(array) == array[np.argsort(array)] should be True. I have found out that this is indeed correct if np.ndim(array) == 2, but it gives different results if np.ndim(array) > 2.

Example:

>>> array = np.array([[[ 0.81774634,  0.62078744],                        [ 0.43912609,  0.29718462]],                       [[ 0.1266578 ,  0.82282054],                        [ 0.98180375,  0.79134389]]]) >>> np.sort(array) array([[[ 0.62078744,  0.81774634],         [ 0.29718462,  0.43912609]],         [[ 0.1266578 ,  0.82282054],         [ 0.79134389,  0.98180375]]]) >>> array.argsort() array([[[1, 0],         [1, 0]],         [[0, 1],         [1, 0]]]) >>> array[array.argsort()] array([[[[[ 0.1266578 ,  0.82282054],           [ 0.98180375,  0.79134389]],           [[ 0.81774634,  0.62078744],           [ 0.43912609,  0.29718462]]],           [[[ 0.1266578 ,  0.82282054],           [ 0.98180375,  0.79134389]],           [[ 0.81774634,  0.62078744],           [ 0.43912609,  0.29718462]]]],           [[[[ 0.81774634,  0.62078744],           [ 0.43912609,  0.29718462]],           [[ 0.1266578 ,  0.82282054],           [ 0.98180375,  0.79134389]]],           [[[ 0.1266578 ,  0.82282054],           [ 0.98180375,  0.79134389]],           [[ 0.81774634,  0.62078744],           [ 0.43912609,  0.29718462]]]]]) 

So, can anybody explain to me how exactly np.argsort() can be used as the indices to obtain the sorted array? The only way I can come up with is:

args = np.argsort(array) array_sort = np.zeros_like(array) for i in range(array.shape[0]):     for j in range(array.shape[1]):         array_sort[i, j] = array[i, j, args[i, j]] 

which is extremely tedious and cannot be generalized for any given number of dimensions.

回答1:

Here is a general method:

import numpy as np  array = np.array([[[ 0.81774634,  0.62078744],                    [ 0.43912609,  0.29718462]],                   [[ 0.1266578 ,  0.82282054],                    [ 0.98180375,  0.79134389]]])  a = 1 # or 0 or 2  order = array.argsort(axis=a)  idx = np.ogrid[tuple(map(slice, array.shape))] # if you don't need full ND generality: in 3D this can be written # much more readable as # m, n, k = array.shape # idx = np.ogrid[:m, :n, :k]  idx[a] = order  print(np.all(array[idx] == np.sort(array, axis=a))) 

Output:

True 

Explanation: We must specify for each element of the output array the complete index of the corresponding element of the input array. Thus each index into the input array has the same shape as the output array or must be broadcastable to that shape.

The indices for the axes along which we do not sort/argsort stay in place. We therefore need to pass a broadcastable range(array.shape[i]) for each of those. The easiest way is to use ogrid to create such a range for all dimensions (If we used this directly, the array would come back unchanged.) and then replace the index correspondingg to the sort axis with the output of argsort.

回答2:

Here's a vectorized implementation. It should be N-dimensional and quite a bit faster than what you're doing.

import numpy as np   def sort1(array, args):     array_sort = np.zeros_like(array)     for i in range(array.shape[0]):         for j in range(array.shape[1]):             array_sort[i, j] = array[i, j, args[i, j]]      return array_sort   def sort2(array, args):     shape = array.shape     idx = np.ix_(*tuple(np.arange(l) for l in shape[:-1]))     idx = tuple(ar[..., None] for ar in idx)     array_sorted = array[idx + (args,)]      return array_sorted   if __name__ == '__main__':     array = np.random.rand(5, 6, 7)     idx = np.argsort(array)      result1 = sort1(array, idx)     result2 = sort2(array, idx)      print(np.array_equal(result1, result2)) 

回答3:

@Hameer's answer works, though it might use some simplification and explanation.

sort and argsort are working on the last axis. argsort returns a 3d array, same shape as the original. The values are the indices on that last axis.

In [17]: np.argsort(arr, axis=2) Out[17]:  array([[[1, 0],         [1, 0]],         [[0, 1],         [1, 0]]], dtype=int32) In [18]: _.shape Out[18]: (2, 2, 2) In [19]: idx=np.argsort(arr, axis=2) 

To use this we need to construct indices for the other dimensions that broadcast to the same (2,2,2) shape. ix_ is a handy tool for this.

Just using idx as one of the ix_ inputs doesn't work:

In [20]: np.ix_(range(2),range(2),idx) .... ValueError: Cross index must be 1 dimensional 

Instead I use the last range, and then ignore it. @Hameer instead constructs the 2d ix_, and then expands them.

In [21]: I,J,K=np.ix_(range(2),range(2),range(2)) In [22]: arr[I,J,idx] Out[22]:  array([[[ 0.62078744,  0.81774634],         [ 0.29718462,  0.43912609]],         [[ 0.1266578 ,  0.82282054],         [ 0.79134389,  0.98180375]]]) 

So the indices for the other dimensions work with the (2,2,2) idx array:

In [24]: I.shape Out[24]: (2, 1, 1) In [25]: J.shape Out[25]: (1, 2, 1) 

That's the basics for constructing the other indices when you are given multidimensional index for one dimension.

@Paul constructs the same indices with ogrid:

In [26]: np.ogrid[slice(2),slice(2),slice(2)]  # np.ogrid[:2,:2,:2] Out[26]:  [array([[[0]],          [[1]]]), array([[[0],          [1]]]), array([[[0, 1]]])] In [27]: _[0].shape Out[27]: (2, 1, 1) 

ogrid as a class works with slices, while ix_ requires a list/array/range.

argsort for a multidimensional ndarray (from 2015) works with a 2d array, but the same logic applies (find a range index(s) that broadcasts with the argsort).

文章来源: How to use numpy.argsort() as indices in more than 2 dimensions?