Comment obtenir des indices de N valeurs maximales dans un tableau numpy?

Question

Numpy propose un moyen d'obtenir l'indice de la valeur maximale d'un tableau via np.argmax .

Je voudrais une chose similaire, mais en retournant les index des N valeurs maximales.

Par exemple, si j'ai un tableau [1, 3, 2, 4, 5] , il function(array, n=3) retournerait [4, 3, 1] .

Merci :)

Answer 1

Nouveaux NumPy versions (1.8) ont une fonction appelée argpartition pour cette. Pour obtenir les indices des quatre grands éléments, ne

>>> a
array([9, 4, 4, 3, 3, 9, 0, 4, 6, 0])
>>> ind = np.argpartition(a, -4)[-4:]
>>> ind
array([1, 5, 8, 0])
>>> a[ind]
array([4, 9, 6, 9])

Contrairement aux argsort, cette fonction s'exécute en temps linéaire dans le pire des cas, mais le retour sur les indices ne sont pas triés, comme on peut le voir à partir du résultat de l'évaluation de l' a[ind]. Si vous avez besoin que trop, de les trier par la suite:

>>> ind[np.argsort(a[ind])]
array([1, 8, 5, 0])

Pour obtenir le top-k éléments dans l'ordre de tri dans cette voie prend O(n + k lg k) temps.

Answer 2

Le plus simple que j'ai pu trouver est:

 In [1]: import numpy as np

In [2]: arr = np.array([1, 3, 2, 4, 5])

In [3]: arr.argsort()[-3:][::-1]
Out[3]: array([4, 3, 1])
 

Cela implique une sorte complète du tableau. Je me demande si numpy fournit un moyen intégré de faire un tri partiel; jusqu'à présent, je n'ai pas réussi à en trouver un.

Si cette solution s'avère trop lente (en particulier pour les petits n ), il peut être intéressant de coder quelque chose en Cython .

Answer 3

EDIT: modifié pour inclure l'amélioration de Ashwini Chaudhary.

 >>> import heapq
>>> import numpy
>>> a = numpy.array([1, 3, 2, 4, 5])
>>> heapq.nlargest(3, range(len(a)), a.take)
[4, 3, 1]
 

Pour les listes Python régulières:

 >>> a = [1, 3, 2, 4, 5]
>>> heapq.nlargest(3, range(len(a)), a.__getitem__)
[4, 3, 1]
 

Si vous utilisez Python 2, utilisez xrange au lieu de range .

Source: http://docs.python.org/3/library/heapq.html

Answer 4

Cela sera plus rapide qu'un tri complet en fonction de la taille de votre tableau d'origine et de la taille de votre sélection:

 >>> A = np.random.randint(0,10,10)
>>> A
array([5, 1, 5, 5, 2, 3, 2, 4, 1, 0])
>>> B = np.zeros(3, int)
>>> for i in xrange(3):
...     idx = np.argmax(A)
...     B[i]=idx; A[idx]=0 #something smaller than A.min()
...     
>>> B
array([0, 2, 3])
 

Bien entendu, cela implique d’altérer votre tableau d’origine. Ce que vous pourriez corriger (si nécessaire) en faisant une copie ou en remplaçant les valeurs d'origine. ... selon ce qui est le moins cher pour votre cas d'utilisation.

Answer 5

bottleneck a une fonction de tri partiel, si le coût du tri du tableau entier pour obtenir les N valeurs les plus grandes est trop important.

_{Je ne connais rien à ce module; Je viens de googler numpy partial sort .}