Have numpy argsort return an array of 2d indices?

Question

If we have a 1d array

arr = np.random.randint(7, size=(5))
# [3 1 4 6 2]
print np.argsort(arr)
# [1 4 0 2 3] <= The indices in the sorted order

If we have a 2d array

arr = np.random.randint(7, size=(3, 3))
# [[5 2 4]
# [3 3 3]
# [6 1 2]]
print np.argsort(arr)
# [[1 2 0]
# [0 1 2]
# [1 2 0]] <= It sorts each row

What I need is the 2d indices that sort this matrix in its entirety. Something like this:

# [[2 1] => 1
# [0 1] => 2
# [2 2] => 2
# .
# .
# .
# [0 2] => 4
# [0 0] => 5
# [2 0]] => 6

How do I get "2d indices" for the sorting of a 2d array?

score 55 · Accepted Answer · answered Jun 01 '15 at 15:34

55

Apply numpy.argsort on flattened array and then unravel the indices back to (3, 3) shape:

>>> arr = np.array([[5, 2, 4],
[3, 3, 3],
[6, 1, 2]])
>>> np.dstack(np.unravel_index(np.argsort(arr.ravel()), (3, 3)))
array([[[2, 1],
        [0, 1],
        [2, 2],
        [1, 0],
        [1, 1],
        [1, 2],
        [0, 2],
        [0, 0],
        [2, 0]]])

answered Jun 01 '15 at 15:34

Ashwini Chaudhary

244,495
58
464
504

5

jezz, what is this result...? – MartianMartian Apr 04 '17 at 16:52
2

I think you need a `[0]` at the end? – Dirk Groeneveld Nov 05 '19 at 23:51
1

this looks more complex than it should be. – pietz Dec 17 '20 at 11:05

Mozak · Answer 2 · 2021-04-12T13:04:57.777

11

From the documentation on numpy.argsort:

ind = np.unravel_index(np.argsort(x, axis=None), x.shape)

Indices of the sorted elements of a N-dimensional array.

An example:

>>> x = np.array([[0, 3], [2, 2]])
>>> x
array([[0, 3],
       [2, 2]])
>>> ind = np.unravel_index(np.argsort(x, axis=None), x.shape)
>>> ind # a tuple of arrays containing the indexes
(array([0, 1, 1, 0]), array([0, 0, 1, 1]))
>>> x[ind]  # same as np.sort(x, axis=None)
array([0, 2, 2, 3])

edited Apr 12 '21 at 13:04

answered Oct 13 '20 at 15:58

Mozak

315
4
7

What edits would have to be made in order to remove the indices corresponding to `np.nan` values in `x` from `ind`? – Ouroboroski Jun 28 '21 at 18:04

score 3 · Answer 3 · answered Feb 08 '22 at 03:55

3

Use np.take_along_axis to slices with your indices. Assuming you need to sort row-wise:

>>> arr = np.array([[5, 2, 4],
[3, 3, 3],
[6, 1, 2]])
>>> ind = arr.argsort(axis=-1)
>>> ind
array([[1, 2, 0],
       [0, 1, 2],
       [1, 2, 0]])
>>> np.take_along_axis(arr, ind, axis=-1)
array([[2, 4, 5],
       [3, 3, 3],
       [1, 2, 6]])

answered Feb 08 '22 at 03:55

Ivelin

12,293
5
37
35

This did not sort all the elements, it sorts rows. – AturSams Feb 08 '22 at 12:21

Have numpy argsort return an array of 2d indices?

3 Answers3

Linked