numpy配列からｎ番目に小さい値まで取得

結構よく出会うシチュエーションですがいつもどうするか迷うのでメモします．timeitで時間を測りましたが参考程度です．

環境

Ubuntu 20.04 on win
python 3.10.4
numpy 1.21.6

準備

適当な配列を用意する．

np.random.seed(0)
a = np.random.randn(1000, 100, 100)

sort + partition

コメントで教えてもらった方法．partitionは小さい方からn番目のもので仕切る．無駄が少ない．

n = 1000
a_minima = np.sort(np.partition(a.ravel(), n)[:n]); a_minima

timeit出力

159 ms ± 18 ms per loop (mean ± std. dev. of 3 runs, 4 loops each)

argsort

インデックスが必要な場合．丸ごとソートするので無駄が多いのと，配列が大きいと終わらない．

n = 1000
idx = np.unravel_index( np.argsort(a=a, axis=None), a.shape)
a_minima = a[idx][:n]

timeit出力

3.47 s ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)

forループ＋argmin

原始的な方法．ｎが小さくないと使えない．

n = 1000
a_minima = []
for _ in range(n):
    idx = np.unravel_index( np.argmin(a=a, axis=None), a.shape )
    a_minima.append( a[idx] )
    a[idx] = np.Inf

timeit出力

10.1 s ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)

argwhere＋argsort

現実には配列の取りうる値の範囲に見当がつくことも多いので，適当なｘについてそれより小さい要素のみ抜き出し，配列サイズを小さくしてから上記の処理を行う．argwhereのアウトプットをそのままインデックスとして使えないのが歯がゆい．

threshold = ( np.max(a) - np.min(a) ) / 5 + np.min(a) # 小さい方から２０％
idx = np.transpose( np.argwhere(a < threshold) )
a_minima = a[idx[0], idx[1], idx[2]]
a_minima_sorted = a_minima[ np.argsort(a_minima) ]

timeit出力

6.8 ms ± 13.5 ms per loop (mean ± std. dev. of 3 runs, 3 loops each)

※追記訂正：argwhereではなくnonzeroを使って同じことができ，かつ返り値をそのままインデックス配列として使える．インデックスが不要ならargsortではなくsortで．

idx = np.nonzero(a < threshold)
idx_to_sort = np.argsort(a[idx])
a_minima = a[idx_to_sort]

最小値とそのインデックスだけ知りたい

最小値だけでよいなら以下のように．

idx = np.unravel_index( np.argmin(a), a.shape)
a_min = a[idx]