How do I get indices of N maximum values in a NumPy array?
Как мне получить индексы из N максимальных значений в массиве NumPy?
NumPy предлагает способ получить индекс максимального значения массива через np.argmax.
Я хотел бы нечто подобное, но возвращающее индексы из N максимальных значений.
Например, если у меня есть массив, [1, 3, 2, 4, 5]тогда nargmax(array, n=3) возвращались бы индексы, [4, 3, 1] которые соответствуют элементам [5, 4, 3].
Переведено автоматически
Ответ 1
В более новых версиях NumPy (1.8 и выше) для этого вызывается функция argpartition. Чтобы получить индексы четырех самых больших элементов, выполните
>>> 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) времени.
Это предполагает полную сортировку массива. Интересно, numpy предоставляет ли встроенный способ выполнения частичной сортировки; пока я не смог его найти.
Если это решение окажется слишком медленным (особенно для небольших n), возможно, стоит подумать о кодировании чего-либо в Cython.