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?
Answers:
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]
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]
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?
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]
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]