Un moyen rapide pour trouver le plus grand de N éléments dans un tableau numpy

numpy 1.8 implémente partition et argpartition qui effectuent partielle de tri ( en O(n) fois, par opposition à une sorte qui est O(n) * log(n)).

import numpy as np

test = np.array([9,1,3,4,8,7,2,5,6,0])

temp = np.argpartition(-test, 4)
result_args = temp[:4]

temp = np.partition(-test, 4)
result = -temp[:4]

Résultat:

>>> result_args
array([0, 4, 8, 5]) # indices of highest vals
>>> result
array([9, 8, 6, 7]) # highest vals

Calendrier:

In [16]: a = np.arange(10000)

In [17]: np.random.shuffle(a)

In [18]: %timeit np.argsort(a)
1000 loops, best of 3: 1.02 ms per loop

In [19]: %timeit np.argpartition(a, 100)
10000 loops, best of 3: 139 us per loop

In [20]: %timeit np.argpartition(a, 1000)
10000 loops, best of 3: 141 us per loop

J'ai eu ce problème et, depuis que cette question est de 5 ans, j'ai dû refaire tous les indices de référence et de modifier la syntaxe de goulot d'étranglement (il n'y a pas de partsort plus, c'est partition maintenant).

J'ai utilisé les mêmes arguments que kwgoodman, sauf le nombre d'éléments récupérés, dont j'ai augmenté à 50 (pour mieux l'adapter à ma situation particulière).

J'ai obtenu ces résultats:

bottleneck 1: 01.12 ms per loop
bottleneck 2: 00.95 ms per loop
pandas      : 01.65 ms per loop
heapq       : 08.61 ms per loop
numpy       : 12.37 ms per loop
numpy 2     : 00.95 ms per loop

Donc, bottleneck_2 et numpy_2 (adas de la solution) étaient attachés.
Mais, à l'aide de np.percentile (numpy_2) vous avez ces topN éléments déjà triés, ce qui n'est pas le cas pour les autres solutions. D'autre part, si vous êtes également intéressé à l'index de ces éléments, centile n'est pas utile.

J'ai ajouté les pandas trop, qui utilise goulot d'étranglement en dessous, si disponible (http://pandas.pydata.org/pandas-docs/stable/install.html#recommended-dependencies). Si vous avez déjà une pandas Série ou DataFrame pour commencer, vous êtes en de bonnes mains, il suffit d'utiliser nlargest et vous avez terminé.

Le code utilisé pour l'indice de référence est comme suit (python 3, s'il vous plaît):

import time
import numpy as np
import bottleneck as bn
import pandas as pd
import heapq
def 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_loop
t1 = 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))

Vous pouvez également utiliser numpy du percentile de la fonction. Dans mon cas, il était légèrement plus rapide goulot d'étranglement.partsort():

import timeit
import bottleneck as bn
N,M,K = 10,1000000,100
start = timeit.default_timer()
for k in range(K):
a=np.random.uniform(size=M)
tmp=-bn.partsort(-a, N)[:N]
stop = timeit.default_timer()
print (stop - start)/K
start = timeit.default_timer()
perc = (np.arange(M-N,M)+1.0)/M*100
for k in range(K):
a=np.random.uniform(size=M)
tmp=np.percentile(a,perc)
stop = timeit.default_timer()
print (stop - start)/K

Temps moyen par boucle:

• goulot d'étranglement.partsort(): 59 ms
• np.centile(): 54 ms