0

The indexes of the k lowest elements of an array can be found using np.argpartition as in this question, I can only think of applying this to many rows using a loop. Is there a way that this can be done without a loop, using some kind of matrix operation? I would like a resulting matrix where each row contains k indexes (of the k lowest elements).

whim
  • 3
  • 3
  • Just use the `axis` argument in the `argpartition` function and it will perform the operation along that axis. – Lith Mar 06 '21 at 18:08

1 Answers1

0

You don't have to worry about having to use loops in NumPy. All such operations are broadcastable. If you want to apply the op to a matrix and have each row containing k lowest elements, use np.argpartition(matrix, axis=-1)

Captain Trojan
  • 1,109
  • 5
  • 11