Problem with understanding of work of np.argpartition-CodePudding

I have problem with execution of np.argpartition I have nd.array

example = np.array([[5,6,7,3,4],[1,2,3,7,5],[6,7,4,2,3],[1,2,3,5,9],[2,3,6,1,2,]])
out: [[5 6 7 3 4]
      [1 2 3 7 5]
      [6 7 4 2 3]
      [1 2 3 5 9]
      [2 3 6 1 2]]

I can get indices for sorted array by np.argsort

print(np.argsort(example))
out:
     [[3 4 0 1 2]
      [0 1 2 4 3]
      [3 4 2 0 1]
      [0 1 2 3 4]
      [3 0 4 1 2]]

I want to use np.argsort to economy some time for executing, because I need only 3 sorted element in each row of this array. I use this code to do it:

print(np.argpartition(example, 3, axis=1))
out: [[3 4 0 1 2]
      [1 0 2 4 3]
      [3 4 2 0 1]
      [1 0 2 3 4]
      [3 4 0 1 2]]

I expect that the first three indices of each row will match the indices in the sorted array, but this is not the caseю That doesn't work . I don't understand what I did wrong.

CodePudding user response：

np.argpartition(example, k, axis=1) does not return sorted array for first k elements. It only returns indices such that only (k 1)th element is sorted. If you see in your output, only the 4th element matches with argsort()

If you want first three sorted elements, you have to give a list for k parameter

index_array = np.argpartition(example, [0,1,2], axis=1)
print(np.take_along_axis(example,index_array, axis=1)) ##this will give you first 3 sorted elements