Numpy：从 nD 数组中提取最小值的索引

Question

I have an array

X = [[0.65108716 0.72213542 0.62142414 0.80734795 0.79485172 0.83946013
0.79192978 0.76614672 0.         0.84231442]
 [0.71353155 0.58493483 0.76903558 0.77678972 0.71837986 0.56127471
0.72591233 0.75986564 0.83495295 0.03016315]]

I need to extract indices of the top k min values from each subarray.

If I want to get only 1 value, I can use

top_n_indices = np.argsort(X)[:, :1]

It gives me indices

[[8], [9]]

And then when I try to extract values with

np.take(X, top_n_indices)

It gives me incorrect answer

[[0.        ], [0.84231442]]

But it should be

[[0.        ], [0.03016315]]

Is it possible to do that without list comprehension?

Answer 1

At the end of the day, do you want the k minimum values from each subarray, or do you need their positions as well? If you only need the values, try the following (eg, to get the 3 lowest values in each subarray):

import numpy as np

X = np.array([[0.65108716,0.72213542,0.62142414,0.80734795,0.79485172,0.83946013,0.79192978,0.76614672,0.,0.84231442],
[0.71353155,0.58493483,0.76903558,0.77678972,0.71837986,0.56127471,0.72591233,0.75986564,0.83495295,0.03016315]])

k=3
k_min_values=np.sort(X,axis=1)[:,:k]

print(k_min_values)

which will yield:

[[0.         0.62142414 0.65108716]
 [0.03016315 0.56127471 0.58493483]]