NumPyは配列の最大値のインデックスを取得する方法を提案しています。np.argmax
。
同様のことをしたいのですが、最大値のインデックスを返しますN
。
たとえば、配列 がある場合[1, 3, 2, 4, 5]
、 は要素 に対応するnargmax(array, n=3)
インデックスを返します。[4, 3, 1]
[5, 4, 3]
ベストアンサー1
新しいNumPyバージョン(1.8以降)には、argpartition
4つの最大の要素のインデックスを取得するには、次のようにします。
>>> a = np.array([9, 4, 4, 3, 3, 9, 0, 4, 6, 0])
>>> a
array([9, 4, 4, 3, 3, 9, 0, 4, 6, 0])
>>> ind = np.argpartition(a, -4)[-4:]
>>> ind
array([1, 5, 8, 0])
>>> top4 = a[ind]
>>> top4
array([4, 9, 6, 9])
とは異なりargsort
、この関数は最悪の場合でも線形時間で実行されますが、 の評価結果からわかるように、返されるインデックスはソートされていませんa[ind]
。ソートも必要な場合は、後で並べ替えます。
>>> ind[np.argsort(a[ind])]
array([1, 8, 5, 0])
この方法でソートされた順序で上位k要素を取得するには、O( n + k log k ) 時間がかかります。