Найти k наименьших значений массива numpy

Чтобы найти индекс наименьшего значения, я могу использовать 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

Но как я могу найти индексы k-наименьших значений?

Я ищу что-то вроде:

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 наименьших значений. k никогда не будет > 10.

Ответы

Ответ 1

Используйте 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[-3:]]
# [  9.  17.  17.]

Вот сравнение с 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

Ответ 2

Вы можете использовать numpy.argsort с нарезкой

>>> import numpy as np
>>> A = np.array([1, 7, 9, 2, 0.1, 17, 17, 1.5])
>>> np.argsort(A)[:3]
array([4, 0, 7], dtype=int32)

Ответ 3

numpy.partition(your_array, k) является альтернативой. Нет необходимости в нарезке, поскольку он задает значения, отсортированные до элемента kth.