2

Related to How to get indices of N maximum values in a numpy array?, I have a numpy matrix a, and I would like to produce the array whose i-th row is the column indices of the top N elements of the i-th row of a.

Following the top-voted answer to the linked question, adapting it for arrays, here is what I have so far (using N=4):

>>> a
array([[9, 4, 4, 3, 3, 9, 0, 4, 6, 0],
       [3, 4, 6, 9, 5, 7, 1, 2, 8, 4]])
>>> ind=np.argpartition(a,-4)[:,-4:]
>>> ind
array([[1, 5, 8, 0],
       [2, 3, 8, 5]])
>>> rows=np.transpose([np.arange(a.shape[0])])
>>> rows
array([[0],
       [1]])
>>> ind_sorted = ind[rows,np.argsort(a[rows,ind])]
>>> ind_sorted
array([[1, 8, 5, 0],
       [2, 5, 8, 3]])

This works, but seems to be not very (num)pythonic. I'm sure there's a better way to do the indexing that doesn't require a dummy array. Any suggestions?

Community
  • 1
  • 1
Matthew Leingang
  • 2,579
  • 3
  • 20
  • 19

2 Answers2

4

Slicing the last four elements of the order index by row seems to be working:

a.argsort(axis = 1)[:, -4:]

# array([[7, 8, 0, 5],
#        [2, 5, 8, 3]])

The tie method is not defined, so there will be some difference between 1 and 7 as well as the order of 0 and 5.

Psidom
  • 209,562
  • 33
  • 339
  • 356
  • Thanks for this answer. I had wanted the `argpartition` method because I was worried about the time to sort the whole matrix (it's about 150x150). Should I be, or is that premature optimization? – Matthew Leingang Jul 08 '16 at 13:55
  • If your matrix is of 150 * 150 dimension, I would not worry about the time. With that being said, the sorting algorithm should be optimized for any high level language if provided, I believe. – Psidom Jul 08 '16 at 14:00
0

Here I am posting code for getting indices for top 2 values in each row.

Unsorted_Array = np.array ([[23, 56, 12, 24], [6, 36, 9, 99], [24, 11, 87, 9], [7,29,103, 5]])
Index_of_Top_2_Values_of_EachRow = np.argsort(Unsorted_Array,axis = 1)[:,-2:]
print(Index_of_Top_2_Values_of_EachRow)
David Buck
  • 3,752
  • 35
  • 31
  • 35
Ravi K
  • 1
  • Thank you for contributing. Other than changing 4 to 2, can you explain how this answer is different than the one posted by Psidom four years ago? – Matthew Leingang Apr 27 '20 at 12:30