Поиск индекса N самых больших элементов в массиве Python/Список эффективно

Я сожалею заранее, если это дублированный вопрос, я искал эту информацию, но все еще не мог ее найти.

Можно ли организовать массив numpy (или список python), используя очень эффективные индексы N самых больших элементов в порядке убывания?

Например, массив:

a = array([4, 1, 0, 8, 5, 2])

Индексы крупнейших элементов в порядке убывания дадут (учитывая N = 6, все элементы включены):

8 → 3

5 → 4

4 → 0

2 → 5

1 → 1

0 → 2

result = [3, 4, 0, 5, 1, 2]

Я знаю, как сделать это, используя несколько глупый подход (например, сортировку массива и поиск каждого из N номеров для своих индексов), но мне было интересно, есть ли какая-либо эффективная библиотека, например, узкое место или heapq или, возможно, pythonic подход, чтобы сделать это очень быстро. Я должен применять его в нескольких массивах с элементами 300 тыс., Поэтому проблема в производительности.

Спасибо заранее!

UPDATE

Я прочитал ответы и решил использовать их с использованием 300k случайных целых чисел, вот результаты:

решение 1: sorted(range(len(a)), key=lambda i:a[i]) время: 230 мс

решение 2: heapq.nlargest(len(a), zip(a, itertools.count())) время: 396 мс

решение 3: heapq.nlargest(len(a), enumerate(a), key=operator.itemgetter(1)) время: 864 мс

решение 4: def f(a,N): return np.argsort(a)[::-1][:N] (N = len(a)) время: 104 мс

Большое спасибо за быстрые и очень хорошие ответы!

Ответ 1

Вы просмотрели встроенный метод numpy argsort?:

http://docs.scipy.org/doc/numpy/reference/generated/numpy.argsort.html

Я могу сортировать массив с 300 000 случайных поплавков примерно на 29 мс на моей машине, используя этот метод.

def f(a,N):
    return np.argsort(a)[::-1][:N]

Ответ 2

L = [4, 1, 0, 8, 5, 2]
sorted(range(len(L)), key=lambda i:L[i])

Ответ 3

Вы можете использовать heapq, чтобы сделать это достаточно легко:

>>> heapq.nlargest(3, zip(a, itertools.count()))
[(8, 3), (5, 4), (4, 5)]

Кортежи сортируются путем сортировки по первому значению, затем второго и т.д. Это означает, что мы можем просто сделать кортеж (value, index) и отсортировать, указав нам индексы значений (значения также указаны, но мы можем легко выбросить их).

Я использую zip() и itertools.count(), поскольку перечисление дает нам неправильный порядок, поэтому они будут отсортированы по индексу, а не по значению. В качестве альтернативы вы также можете сделать ((value, index) for index, value in enumerate(a)), но я чувствую, что это менее понятно.

Другой вариант - дать ключ, сделав heapq.nlargest(3, enumerate(a), key=operator.itemgetter(1)).

Ответ 4

Другой способ использования heapq

heapq.nlargest(n, range(len(a)), key=a.__getitem__)

Как прокомментировано в другом месте, он не будет бить сортировку, если он не очень большой и n<<len(a), потому что сортировка является относительно быстрой операцией в Python. Однако в конечном итоге медленный O (n) всегда будет бить O (n * log (n))