argsort for a multidimensional ndarray
Solution:
>>> a[np.arange(np.shape(a)[0])[:,np.newaxis], np.argsort(a)]array([[1, 2, 3], [2, 8, 9]])
You got it right, though I wouldn't describe it as cheating the indexing.
Maybe this will help make it clearer:
In [544]: i=np.argsort(a,axis=1)In [545]: iOut[545]: array([[1, 2, 0], [2, 0, 1]])
i
is the order that we want, for each row. That is:
In [546]: a[0, i[0,:]]Out[546]: array([1, 2, 3])In [547]: a[1, i[1,:]]Out[547]: array([2, 8, 9])
To do both indexing steps at once, we have to use a 'column' index for the 1st dimension.
In [548]: a[[[0],[1]],i]Out[548]: array([[1, 2, 3], [2, 8, 9]])
Another array that could be paired with i
is:
In [560]: j=np.array([[0,0,0],[1,1,1]])In [561]: jOut[561]: array([[0, 0, 0], [1, 1, 1]])In [562]: a[j,i]Out[562]: array([[1, 2, 3], [2, 8, 9]])
If i
identifies the column for each element, then j
specifies the row for each element. The [[0],[1]]
column array works just as well because it can be broadcasted against i
.
I think of
np.array([[0], [1]])
as 'short hand' for j
. Together they define the source row and column of each element of the new array. They work together, not sequentially.
The full mapping from a
to the new array is:
[a[0,1] a[0,2] a[0,0] a[1,2] a[1,0] a[1,1]]
def foo(a): i = np.argsort(a, axis=1) return (np.arange(a.shape[0])[:,None], i)In [61]: foo(a)Out[61]: (array([[0], [1]]), array([[1, 2, 0], [2, 0, 1]], dtype=int32))In [62]: a[foo(a)]Out[62]: array([[1, 2, 3], [2, 8, 9]])
I found the answer here, with someone having the same problem. They key is just cheating the indexing to work properly...
>>> a[np.arange(np.shape(a)[0])[:,np.newaxis], np.argsort(a)]array([[1, 2, 3], [2, 8, 9]])
The above answers are now a bit outdated, since new functionality was added in numpy 1.15 to make it simpler; take_along_axis (https://docs.scipy.org/doc/numpy-1.15.1/reference/generated/numpy.take_along_axis.html) allows you to do:
>>> a = np.array([[3,1,2],[8,9,2]])>>> np.take_along_axis(a, a.argsort(axis=-1), axis=-1)array([[1 2 3] [2 8 9]])