Is there a numpy limited argsort?

Viewed 62

I have an array with 1M floats and want to get the top 10. Wondering if there's an argsort that will only give top 10 to speed things up. I suppose I could do np.argmax() 10 times, dropping the resulting index each time (probably by making it -inf, to not have to keep track of shifting indices). Just checking if there's a well-known inbuilt way.

Doesn't necessarily have to be numpy. Second best is another very common library.

1 Answers

You can indeed use numpy.argpartition():

import numpy
a = numpy.random.rand(1000000)
b = numpy.argpartition(a, [0,1,2,3,4,5,6,7,8,9])[:10]

It's nearly 10 times faster (in this case, on my machine):

>>> import timeit, numpy
>>> a = numpy.random.rand(1000000)
>>> # Sorting traditionally
>>> timeit.timeit('numpy.argsort(a)[:10]', globals=globals(), number=100)
10.2261...
>>> # Getting the top 10 but not in order
>>> timeit.timeit('numpy.argpartition(a, 10)[:10]', globals=globals(), number=100)
1.2125...
>>> # Getting the top 10 in order
>>> timeit.timeit('numpy.argpartition(a, [0,1,2,3,4,5,6,7,8,9])[:10]', globals=globals(), number=100)
1.4764...
Related