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)).

Answer 2

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

Answer 3

Another way to use heapq

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

As commented elsewhere, it won't beat sorting unless a is very large and n<<len(a) because sorting is a relatively fast operation in Python. However eventually a slow O(n) will always beat the O(n*log(n))

Answer 4

Have you looked at the built-in numpy argsort method?:

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

I can sort an array with 300,000 random floats in about 29 ms on my machine using that method.

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