Sorting 2D numpy array using indices returned from np.argsort() [duplicate]

Question

When we have a 1D numpy array, we can sort it the following way:

>>> temp = np.random.randint(1,10, 10)
>>> temp
array([5, 1, 1, 9, 5, 2, 8, 7, 3, 9])
>>> sort_inds = np.argsort(temp)
>>> sort_inds
array([1, 2, 5, 8, 0, 4, 7, 6, 3, 9], dtype=int64)
>>> temp[sort_inds]
array([1, 1, 2, 3, 5, 5, 7, 8, 9, 9])

Note: I know I can do this using np.sort; Obviously, I need the sorting indices for a different array - this is just a simple example. Now we can continue to my actual question..

I tried to apply the same approach for a 2D array:

>>> d = np.random.randint(1,10,(5,10))
>>> d
array([[1, 6, 8, 4, 4, 4, 4, 4, 4, 8],
       [3, 6, 1, 4, 5, 5, 2, 1, 8, 2],
       [1, 2, 6, 9, 8, 6, 9, 2, 5, 8],
       [8, 5, 1, 6, 6, 2, 4, 3, 7, 1],
       [5, 1, 4, 4, 4, 2, 5, 9, 7, 9]])
>>> sort_inds = np.argsort(d)
>>> sort_inds
array([[0, 3, 4, 5, 6, 7, 8, 1, 2, 9],
       [2, 7, 6, 9, 0, 3, 4, 5, 1, 8],
       [0, 1, 7, 8, 2, 5, 4, 9, 3, 6],
       [2, 9, 5, 7, 6, 1, 3, 4, 8, 0],
       [1, 5, 2, 3, 4, 0, 6, 8, 7, 9]], dtype=int64)

This result looks good - notice that we can sort each row of d using the indices of the corresponding row from sort_inds as demonstrated in the 1D example. However, trying to get a sorted array using the same approach I used in the 1D example, I got this exception:

>>> d[sort_inds]
---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
<ipython-input-63-e480a9fb309c> in <module>
----> 1 d[ind]

IndexError: index 5 is out of bounds for axis 0 with size 5

So I have 2 questions:

1. What just happened? How did numpy interpret this code?
2. How can I still achieve what I want - that is, sorting d - or any other array of the same dimensions - using sort_inds?

Thanks

Answer 1

You need a little extra work to properly index the 2d array. Here's a way using advanced indexing, where np.arange is used in the first axis so that each row in sort_inds extracts values from the corresponding row in d:

d[np.arange(d.shape[0])[:,None], sort_inds]

array([[1, 1, 2, 3, 3, 4, 4, 7, 8, 9],
       [1, 3, 4, 5, 5, 5, 6, 8, 8, 9],
       [1, 2, 3, 4, 5, 6, 7, 8, 8, 8],
       [2, 2, 4, 7, 7, 8, 8, 9, 9, 9],
       [1, 1, 2, 4, 4, 7, 7, 8, 8, 8]])