NumPy: argsort to give row-only permutation?

Question:

I can use np.argsort() to create a permutation which if used as an indexing expression returns elements in order:

>>> a = np.array([9, 3, 5, 2, 5, 1, 0, 8])
>>> idx = np.argsort(a)
>>> a[idx]
array([0, 1, 2, 3, 5, 5, 8, 9])

All well and good. Now I want to use this technique on a matrix, to give a permutation which if used for indexing would select the rows in order. Note that I don’t want the results of applying argsort() on axis 0, as this gives a permutation per column.

In other words,

>>> a = np.array([7, 3, 4, 7, 9, 7, 3, 4, 5]).reshape((3,3))
>>> a
array([[7, 3, 4],
       [7, 9, 7],
       [3, 4, 5]])
>>> np.find_lexicographic_row_permutation(a) # pseudo-code
array([2, 0, 1])

In other words, the partition would reorder the matrix’s rows in ‘lexicographic’ order.

How can I achieve this in a generic way?

Asked By: xpqz

||

Answers:

You can use np.lexsort to solve that:

row_permutation = np.lexsort(a.T[::-1])

We use the transpose because np.lexsort() sorts the data by columns, and you want to sort the rows lexicographically.

Answered By: Hamzah
Categories: questions Tags: ,
Answers are sorted by their score. The answer accepted by the question owner as the best is marked with
at the top-right corner.