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?
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.
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?
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.