Python numpy keep a list of indices of a sorted 2D array

Question

I have a 2D numpy array and I want to create a new 1D array where it is indices of numbers in the first array if they are sorted in an ascending order. For the following array:

A = [[1,0,2],
     [0,3,0]]

I want this to be like:

B = [[1,1],[0,2],[0,0],[0,1],[1,0],[1,2]]

Any idea how it can be done in python using predefined functions?

Thanks

behzad.nouri · Accepted Answer · 2015-04-19 19:42:02Z

13

You can use argsort to sort the indices of flattened array, followed by unravel_index to convert the flat index back to coordinates:

>>> i = (-a).argsort(axis=None, kind='mergesort')
>>> j = np.unravel_index(i, a.shape) 
>>> np.vstack(j).T
array([[1, 1],
       [0, 2],
       [0, 0],
       [0, 1],
       [1, 0],
       [1, 2]])

-a and kind='mergesort' is in order to sort the array in a stable manner in descending order (to match the output you are looking for).

If you do not care about having a stable sort, replace the first line with:

>>> i = a.argsort(axis=None)[::-1]

edited Apr 19, 2015 at 19:42

answered Apr 19, 2015 at 19:30

behzad.nouri

78.5k18 gold badges130 silver badges127 bronze badges

Sign up to request clarification or add additional context in comments.

3 Comments

Steven C. Howell Over a year ago

The tip about using np.unravel_index is brilliant! I am glad to learn this method.

Steven C. Howell Over a year ago

What do you mean by "stable sort"? Does this produce consistent sorting for identical values?

ahajib Over a year ago

@stvn66 I believe this does produce consistent sorting for identical values which is equal to stable sorting

Collectives™ on Stack Overflow

Python numpy keep a list of indices of a sorted 2D array

1 Answer 1

3 Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

3 Comments

Your Answer

Sign up or log in

Post as a guest

Related