numpy 배열의 축을 따라 k- 최소값의 인덱스

Question

루프를 사용하지 않고 numpy 배열의 축을 따라 k- 최소값의 인덱스를 반환하는 방법이 있습니까?

Answer 1

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