Quantcast
Channel: A fast way to find the largest N elements in an numpy array - Stack Overflow
Viewing all articles
Browse latest Browse all 8

Answer by Tacio Medeiros for A fast way to find the largest N elements in an numpy array

$
0
0

I had this problem and, since this question is 5 years old, I had to redo all benchmarks and change the syntax of bottleneck (there is no partsort anymore, it's partition now).

I used the same arguments as kwgoodman, except the number of elements retrieved, which I increased to 50 (to better fit my particular situation).

I got these results:

bottleneck 1: 01.12 ms per loopbottleneck 2: 00.95 ms per looppandas      : 01.65 ms per loopheapq       : 08.61 ms per loopnumpy       : 12.37 ms per loopnumpy 2     : 00.95 ms per loop

So, bottleneck_2 and numpy_2 (adas's solution) were tied.But, using np.percentile (numpy_2) you have those topN elements already sorted, which is not the case for the other solutions. On the other hand, if you are also interested on the indexes of those elements, percentile is not useful.

I added pandas too, which uses bottleneck underneath, if available (http://pandas.pydata.org/pandas-docs/stable/install.html#recommended-dependencies). If you already have a pandas Series or DataFrame to start with, you are in good hands, just use nlargest and you're done.

The code used for the benchmark is as follows (python 3, please):

import timeimport numpy as npimport bottleneck as bnimport pandas as pdimport heapqdef bottleneck_1(a, n):    return -bn.partition(-a, n)[:n]def bottleneck_2(a, n):    return bn.partition(a, a.size-n)[-n:]def numpy(a, n):    return a[a.argsort()[-n:]]def numpy_2(a, n):    M = a.shape[0]    perc = (np.arange(M-n,M)+1.0)/M*100    return np.percentile(a,perc)def pandas(a, n):    return pd.Series(a).nlargest(n)def hpq(a, n):    return heapq.nlargest(n, a)def do_nothing(a, n):    return a[:n]def benchmark(func, size=1000000, ntimes=100, topn=50):    t1 = time.time()    for n in range(ntimes):        a = np.random.rand(size)        func(a, topn)    t2 = time.time()    ms_per_loop = 1000000 * (t2 - t1) / size    return ms_per_loopt1 = benchmark(bottleneck_1)t2 = benchmark(bottleneck_2)t3 = benchmark(pandas)t4 = benchmark(hpq)t5 = benchmark(numpy)t6 = benchmark(numpy_2)t0 = benchmark(do_nothing)print("bottleneck 1: {:05.2f} ms per loop".format(t1 - t0))print("bottleneck 2: {:05.2f} ms per loop".format(t2 - t0))print("pandas      : {:05.2f} ms per loop".format(t3 - t0))print("heapq       : {:05.2f} ms per loop".format(t4 - t0))print("numpy       : {:05.2f} ms per loop".format(t5 - t0))print("numpy 2     : {:05.2f} ms per loop".format(t6 - t0))

Viewing all articles
Browse latest Browse all 8

Trending Articles



<script src="https://jsc.adskeeper.com/r/s/rssing.com.1596347.js" async> </script>