Получить индексы N максимальных значений в массиве NumPy без сортировки их?

Question

Мой вопрос очень похож на этот: Как получить индексы N максимальных значений в массиве numpy?

Но я хотел бы получить индексы в том же порядке. Iнайдите их.

Давайте возьмем пример, помеченный в этом вопросе, как правильное решение:

import numpy as np
arr = np.array([1, 3, 2, 4, 5])
arr.argsort()[-3:][::-1]

array([4, 3, 1])

Вместо этого я должен получить следующий результат:

array([1, 3, 4])

Answer 1

Вероятно, это немного зависит от размеров a и k, но часто наиболее быстрым представляется сочетание partition с flatnonzero или where:

>>> a = np.random.random(10000)
>>> k = 5
>>> 
>>> timeit("np.flatnonzero(a >= np.partition(a, len(a) - k)[len(a) - k])", globals=globals(), number=10000)
0.8328661819687113
>>> timeit("np.sort(np.argpartition(a, len(a) - k)[len(a) - k:])", globals=globals(), number=10000)
1.0577796879806556
>>> np.flatnonzero(a >= np.partition(a, len(a) - k)[len(a) - k])
array([2527, 4299, 5531, 6945, 7174])
>>> np.sort(np.argpartition(a, len(a) - k)[len(a) - k:])
array([2527, 4299, 5531, 6945, 7174])

Примечание 1: это подчеркивает значительные издержки производительности при косвенной индексации.

Примечание 2: поскольку мы используем только элемент pivot и отбрасываем фактический раздел, percentile теоретически должен быть как минимум таким же быстрым, но на практике это намного медленнее.

Answer 2

Использование numpy.argpartition():

k = 3
np.argpartition(arr, len(arr) - k)[-k:]

Настройка индекса k на все, что вам нужно.

ПРИМЕЧАНИЕ: возвращенные индексы не гарантируются в«порядок сортировки» - просто все, что за индексом k больше, чем значение в позиции k в отсортированном массиве.

ПРИМЕЧАНИЕ 2: если вам нужно, чтобы возвращаемые индексы были отсортированы сами, затем просто добавьте numpy.sort() к вышеприведенной команде:

np.sort(np.argpartition(arr, len(arr) - k)[-k:])

numpy.argpartition() обеспечивает значительное повышение производительности по сравнению с полным sort, особенно для больших arr.В приведенном выше примере вы выполняете полную сортировку только по выбранным индексам (не всем).