74

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?

Georgy
  • 9,972
  • 7
  • 57
  • 66
hardikudeshi
  • 1,351
  • 5
  • 17
  • 22

2 Answers2

120
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]
David Beauchemin
  • 184
  • 1
  • 10
Fred Foo
  • 342,876
  • 71
  • 713
  • 819
  • 1
    In Python 3, zip returns an iterable object. So, the last one would have slight modification: `list(zip(...))[0]` – Joseph Apr 23 '20 at 04:00
29

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]
Georgy
  • 9,972
  • 7
  • 57
  • 66
aikramer2
  • 1,175
  • 11
  • 9
  • 3
    If you're willing to use `numpy`, it's more efficient to use `argpartition` than `argsort`. See this answer: https://stackoverflow.com/a/23734295/388951 – danvk Aug 02 '18 at 15:49