How would you find the the indexes of the k lowest elements in each row of a matrix with no looping?

Question

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

Answer 1

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)