[解決済み] numpy配列の最小のk個の値のインデックスを求めます。

質問

最小の値のインデックスを求めるには argmin :

import numpy as np
A = np.array([1, 7, 9, 2, 0.1, 17, 17, 1.5])
print A.argmin()     # 4 because A[4] = 0.1

みたいなのを探しています。

print A.argmin(numberofvalues=3)   
# [4, 0, 7]  because A[4] <= A[0] <= A[7] <= all other A[i]

注：私のユースケースでAは〜10 000〜100 000の値を持ち、私はk=10最小の値のインデックスにのみ興味があります。

どのように解決するのですか？

使用方法 np.argpartition . これは配列全体をソートするわけではありません。それは kth 要素がソートされた位置にあることを保証するだけであり、 それより小さい要素はすべてその前に移動します。したがって、最初の k の要素はk番目に小さい要素になります。

import numpy as np

A = np.array([1, 7, 9, 2, 0.1, 17, 17, 1.5])
k = 3

idx = np.argpartition(A, k)
print(idx)
# [4 0 7 3 1 2 6 5]

これは最も小さいk個の値を返す。これらはソート順でない可能性があることに注意してください。

print(A[idx[:k]])
# [ 0.1  1.   1.5]

k-最大値を得るには

idx = np.argpartition(A, -k)
# [4 0 7 3 1 2 6 5]

A[idx[-k:]]
# [  9.  17.  17.]

警告: (再)使用しないでください idx = np.argpartition(A, k); A[idx[-k:]] を使ってはいけません。 それは常に機能するわけではありません。例えば、これらは x :

x = np.array([100, 90, 80, 70, 60, 50, 40, 30, 20, 10, 0])
idx = np.argpartition(x, 3)
x[idx[-3:]]
array([ 70,  80, 100])

以下は np.argsort との比較です。これも動作しますが、結果を得るために配列全体をソートするだけです。

In [2]: x = np.random.randn(100000)

In [3]: %timeit idx0 = np.argsort(x)[:100]
100 loops, best of 3: 8.26 ms per loop

In [4]: %timeit idx1 = np.argpartition(x, 100)[:100]
1000 loops, best of 3: 721 µs per loop

In [5]: np.alltrue(np.sort(np.argsort(x)[:100]) == np.sort(np.argpartition(x, 100)[:100]))
Out[5]: True