В работе с массивами данных на языке Python и библиотекой NumPy часто может встретиться задача поиска не просто максимального элемента, но и определение индексов N наибольших значений. Например, есть массив чисел [1, 3, 2, 4, 5] и требуется найти индексы трех наибольших значений.
NumPy предлагает удобный метод np.argmax(), позволяющий найти индекс максимального значения в массиве. Однако, этот метод не позволяет сразу находить N максимальных значений. Как же быть? Решение есть!
Используем функцию np.argpartition()
NumPy предоставляет функцию np.argpartition(), которая может быть использована для решения этой проблемы. Эта функция выполняет частичную сортировку входного массива по указанному ключу (в данном случае, по числовым значениям) и возвращает массив индексов такой, что элементы меньше ключа помещаются перед элементами, которые больше или равны ключу.
Вот пример кода, который позволяет найти индексы N наибольших значений в массиве:
import numpy as np
def n_largest_indices(arr, n):
return np.argpartition(arr, -n)[-n:]
array = np.array([1, 3, 2, 4, 5])
n = 3
print(n_largest_indices(array, n))
В данном случае, функция n_largest_indices возвращает индексы трех наибольших значений в массиве array. При этом используется отрицательное значение ключа -n в функции np.argpartition(), что означает, что необходимо вернуть индексы N наибольших, а не наименьших значений.
Важно отметить, что функция np.argpartition() не гарантирует возврата индексов в отсортированном порядке. Если требуется получить индексы в порядке убывания значений, то потребуется дополнительная сортировка.
Вот пример кода, который возвращает отсортированные индексы:
def n_largest_indices_sorted(arr, n):
indices = np.argpartition(arr, -n)[-n:]
return indices[np.argsort(-arr[indices])]
print(n_largest_indices_sorted(array, n))
В этом примере используется функция np.argsort(), которая возвращает индексы элементов в отсортированном порядке, а оператор - используется для сортировки в порядке убывания.
Добавить комментарий