Search code examples
pythonarraysnumpyindicesminim

Indices of k-minimum values along an axis of a numpy array


Is there a way to return the indices of k-minimum values along an axis of a numpy array without using loops?


Solution

  • import numpy as np
    x = np.array([[5, 2, 3],[1, 9, 2]]) # example data
    k = 2 # return the indices of the 2 smallest values
    np.argsort(x, axis=1)[:,0:k] # by row
    
    array([[1, 2],
           [0, 2]])