Finding the Index of N biggest elements in Python Array / List Efficiently

Question

I\'m sorry in advance if this is a duplicated question, I looked for this information but still couldn\'t find it.

Is it possible to arrange a numpy array (or python

Answer 1

You can use heapq to do this easily enough:

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

Tuples are sorted by sorting on the first value, then the second, etc... This means that we can simply make a tuple of (value, index) and sort, giving us the indices of the values (the values are also given, but we can easily throw these away).

I am using zip() and itertools.count() as enumerate gives us the wrong order, so they will be sorted by index, rather than by value. Alternatively, you could also do ((value, index) for index, value in enumerate(a)), but I feel that is less clear.

Another alternative is to give a key, doing heapq.nlargest(3, enumerate(a), key=operator.itemgetter(1)).