Get indices of the top N values of a list

Question

I have a list say a = [5,3,1,4,10]. I need to get indices of the top two values of the list, that is for 5 and 10 I would get [0, 4]. Is there a one-liner that Python offers for such a case?

Overbear

When you say 'values', do you only mean ints? floats? numeric? string? arbitrary object? It's better to give an example list that tests mixed types. – Propane

Answer 1

sorted(range(len(a)), key=lambda i: a[i])[-2:]

or

sorted(range(len(a)), key=lambda i: a[i], reverse=True)[:2]

or

import operator

zip(*sorted(enumerate(a), key=operator.itemgetter(1)))[0][-2:]

or (for long lists), consider using heapq.nlargest

zip(*heapq.nlargest(2, enumerate(a), key=operator.itemgetter(1)))[0]

Luckett

In Python 3, zip returns an iterable object. So, the last one would have slight modification: list(zip(...))[0] – Tatianna

Answer 2

Just a NumPy alternative:

import numpy as np

top_2_idx = np.argsort(a)[-2:]
top_2_values = [a[i] for i in top_2_idx]

Babcock

If you're willing to use numpy, it's more efficient to use argpartition than argsort. See this answer: https://mcmap.net/q/53753/-how-do-i-get-indices-of-n-maximum-values-in-a-numpy-array – Oney

Answer 3

As danvk suggested, we can use argpartion method of numpy for good efficiency like this answer.

If we need two top indexes, we can use this sample code:

import numpy as np
# Find two top indexes
ind = np.argpartition(a, -2)[-2:]
# Sort them
ind = ind[np.argsort(a[ind])]

Azal