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

485

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

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

Par exemple, si j'ai un tableau, [1, 3, 2, 4, 5], function(array, n=3)renverrait les indices [4, 3, 1]qui correspondent aux éléments [5, 4, 3].

Alexis Métaireau
la source
4
Votre question n'est pas vraiment bien définie. Par exemple, ce serait d'être pour les indices (vous attendez) array([5, 1, 5, 5, 2, 3, 2, 4, 1, 5])Whit, n= 3? Lequel de toutes les alternatives, comme [0, 2, 3], [0, 2, 9], ...serait la bonne? Veuillez élaborer davantage sur vos besoins spécifiques. Merci
manger
@eat, je ne me soucie pas vraiment de savoir lequel est censé être retourné dans ce cas spécifique. Même s'il semble logique de rendre le premier rencontré, ce n'est pas une exigence pour moi.
Alexis Métaireau
argsortpourrait être une alternative viable si vous ne vous souciez pas de l'ordre des indécis retournés. Voir ma réponse ci-dessous.
bleu

Réponses:

349

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 de tableau. Je me demande si numpyfournit un moyen intégré de faire un tri partiel; jusqu'à présent, je n'ai pas pu en trouver un.

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

NPE
la source
1
La ligne 3 pourrait-elle s'écrire de manière équivalente arr.argsort()[-1:-4:-1]? Je l'ai essayé dans l'interpréteur et cela donne le même résultat, mais je me demande si ce n'est pas cassé par un exemple.
abroekhof
44
@abroekhof Oui, cela devrait être équivalent à n'importe quelle liste ou tableau. Alternativement, cela pourrait être fait sans inversion en utilisant np.argsort(-arr)[:3], que je trouve plus lisible et pertinent.
askewchan
6
que signifie [:: - 1]? @NPE
1a1a11a
@ 1a1a11a cela signifie inverser un tableau (littéralement, prend une copie d'un tableau de min sans contrainte à max sans contrainte dans un ordre inversé)
FizBack
15
arr.argsort()[::-1][:n]est mieux car il renvoie vide pour n=0au lieu du tableau complet
abora
600

Les versions plus récentes de NumPy (1.8 et plus) ont une fonction appelée argpartitionà cet effet. Pour obtenir les indices des quatre plus grands éléments, faites

>>> a = np.array([9, 4, 4, 3, 3, 9, 0, 4, 6, 0])
>>> 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 à argsort, cette fonction s'exécute en temps linéaire dans le pire des cas, mais les indices renvoyés ne sont pas triés, comme le montre le résultat de l'évaluation a[ind]. Si vous en avez également besoin, triez-les ensuite:

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

Pour obtenir le top- k éléments l' ordre de tri de cette manière prend O ( n + k log k ) temps.

Fred Foo
la source
27
@varela argpartitions'exécute en temps linéaire, O (n), en utilisant l' algorithme introsélection . Le tri suivant ne gère que k éléments, de sorte que s'exécute dans O (k log k).
Fred Foo
2
Si quelqu'un se demande comment fonctionne exactement np.argpartitionet son algorithme sœur, np.partitionil y a une explication plus détaillée dans la question liée: stackoverflow.com/questions/10337533/…
Ramon Martinez
7
@FredFoo: pourquoi avez-vous utilisé -4? avez-vous fait cela pour commencer à l'envers? (puisque k étant positif ou négatif, cela fonctionne de la même manière pour moi! il n'imprime que les plus petits nombres en premier!
Rika
2
@LKT utilise a=np.array([9, 4, 4, 3, 3, 9, 0, 4, 6, 0])parce que les listes python normales ne prennent pas en charge l'indexation par listes, contrairement ànp.array
Marawan Okasha
2
@Umangsinghal np.argpartitionprend un axisargument facultatif . Pour trouver les indices des n premières valeurs pour chaque ligne:np.argpartition(a, -n, axis=1)[-n:]
Ralph
48

Plus simple encore:

idx = (-arr).argsort()[:n]

n est le nombre de valeurs maximales.

Ketan
la source
7
Cela peut-il être fait pour un tableau 2D? Sinon, savez-vous comment?
Andrew Hundt
2
@AndrewHundt: utilisez simplement (-arr) .argsort (axe = -1) [:,: n]
MiniQuark
2
similaire serait arr[arr.argsort()[-n:]]au lieu de nier le tableau, prenez simplement une tranche des n derniers éléments
loganjones16
35

Utilisation:

>>> 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 xrangeplutôt que range.

Source: heapq - Algorithme de file d'attente de tas

anishpatel
la source
2
Il n'y a pas besoin d'une boucle du tout ici: heapq.nlargest(3, xrange(len(a)), a.take). Pour les listes Python, nous pouvons utiliser à la .__getitem__place de .take.
Ashwini Chaudhary
Pour les tableaux n dimensions Aen général: heapq.nlargest(3, range(len(A.ravel())), A.ravel().take). (J'espère que cela ne fonctionne que sur les vues, voir aussi ( ravel vs flatten] ( stackoverflow.com/a/28930580/603003 )).
ComFreek
31

Si vous travaillez avec un tableau multidimensionnel, vous devrez aplatir et démêler les indices:

def largest_indices(ary, n):
    """Returns the n largest indices from a numpy array."""
    flat = ary.flatten()
    indices = np.argpartition(flat, -n)[-n:]
    indices = indices[np.argsort(-flat[indices])]
    return np.unravel_index(indices, ary.shape)

Par exemple:

>>> xs = np.sin(np.arange(9)).reshape((3, 3))
>>> xs
array([[ 0.        ,  0.84147098,  0.90929743],
       [ 0.14112001, -0.7568025 , -0.95892427],
       [-0.2794155 ,  0.6569866 ,  0.98935825]])
>>> largest_indices(xs, 3)
(array([2, 0, 0]), array([2, 2, 1]))
>>> xs[largest_indices(xs, 3)]
array([ 0.98935825,  0.90929743,  0.84147098])
danvk
la source
9

Si vous ne vous souciez pas de l' ordre des K-ème éléments les plus grands que vous pouvez utiliser argpartition, qui devraient être plus performants qu'un tri complet argsort.

K = 4 # We want the indices of the four largest values
a = np.array([0, 8, 0, 4, 5, 8, 8, 0, 4, 2])
np.argpartition(a,-K)[-K:]
array([4, 1, 5, 6])

Les crédits vont à cette question .

J'ai effectué quelques tests et cela ressemble à des argpartitionperformances supérieures à argsortmesure que la taille du tableau et la valeur de K augmentent.

bleu
la source
7

Pour les tableaux multidimensionnels, vous pouvez utiliser le axismot - clé afin d'appliquer le partitionnement le long de l'axe attendu.

# For a 2D array
indices = np.argpartition(arr, -N, axis=1)[:, -N:]

Et pour saisir les objets:

x = arr.shape[0]
arr[np.repeat(np.arange(x), N), indices.ravel()].reshape(x, N)

Mais notez que cela ne retournera pas un résultat trié. Dans ce cas, vous pouvez utiliser le np.argsort()long de l'axe prévu:

indices = np.argsort(arr, axis=1)[:, -N:]

# Result
x = arr.shape[0]
arr[np.repeat(np.arange(x), N), indices.ravel()].reshape(x, N)

Voici un exemple:

In [42]: a = np.random.randint(0, 20, (10, 10))

In [44]: a
Out[44]:
array([[ 7, 11, 12,  0,  2,  3,  4, 10,  6, 10],
       [16, 16,  4,  3, 18,  5, 10,  4, 14,  9],
       [ 2,  9, 15, 12, 18,  3, 13, 11,  5, 10],
       [14,  0,  9, 11,  1,  4,  9, 19, 18, 12],
       [ 0, 10,  5, 15,  9, 18,  5,  2, 16, 19],
       [14, 19,  3, 11, 13, 11, 13, 11,  1, 14],
       [ 7, 15, 18,  6,  5, 13,  1,  7,  9, 19],
       [11, 17, 11, 16, 14,  3, 16,  1, 12, 19],
       [ 2,  4, 14,  8,  6,  9, 14,  9,  1,  5],
       [ 1, 10, 15,  0,  1,  9, 18,  2,  2, 12]])

In [45]: np.argpartition(a, np.argmin(a, axis=0))[:, 1:] # 1 is because the first item is the minimum one.
Out[45]:
array([[4, 5, 6, 8, 0, 7, 9, 1, 2],
       [2, 7, 5, 9, 6, 8, 1, 0, 4],
       [5, 8, 1, 9, 7, 3, 6, 2, 4],
       [4, 5, 2, 6, 3, 9, 0, 8, 7],
       [7, 2, 6, 4, 1, 3, 8, 5, 9],
       [2, 3, 5, 7, 6, 4, 0, 9, 1],
       [4, 3, 0, 7, 8, 5, 1, 2, 9],
       [5, 2, 0, 8, 4, 6, 3, 1, 9],
       [0, 1, 9, 4, 3, 7, 5, 2, 6],
       [0, 4, 7, 8, 5, 1, 9, 2, 6]])

In [46]: np.argpartition(a, np.argmin(a, axis=0))[:, -3:]
Out[46]:
array([[9, 1, 2],
       [1, 0, 4],
       [6, 2, 4],
       [0, 8, 7],
       [8, 5, 9],
       [0, 9, 1],
       [1, 2, 9],
       [3, 1, 9],
       [5, 2, 6],
       [9, 2, 6]])

In [89]: a[np.repeat(np.arange(x), 3), ind.ravel()].reshape(x, 3)
Out[89]:
array([[10, 11, 12],
       [16, 16, 18],
       [13, 15, 18],
       [14, 18, 19],
       [16, 18, 19],
       [14, 14, 19],
       [15, 18, 19],
       [16, 17, 19],
       [ 9, 14, 14],
       [12, 15, 18]])
Kasramvd
la source
Je pense que vous pouvez simplifier l'indexation ici en utilisant np.take_along_axis(qui n'existait probablement pas lorsque vous avez répondu à cette question)
Eric
4

Ce sera plus rapide qu'un tri complet selon la taille de votre tableau d'origine et 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])

Cela implique, bien entendu, de falsifier votre baie d'origine. 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.

Paul
la source
FWIW, votre solution ne fournira pas de solution sans ambiguïté dans toutes les situations. OP doit décrire comment gérer ces cas non ambigus. Merci
manger
@eat La question du PO est un peu ambiguë. Une implémentation, cependant, n'est pas vraiment sujette à interprétation. :) L'OP doit simplement se référer à la définition de np.argmax docs.scipy.org/doc/numpy/reference/generated/numpy.argmax.html pour être sûr que cette solution spécifique répond aux exigences. Il est possible que toute solution répondant aux exigences déclarées du PO soit acceptable.
Paul
Eh bien, on pourrait également considérer la mise en œuvre de argmax(.)sans ambiguïté. (À mon humble avis, il essaie de suivre une sorte de logique de court-circuit, mais ne parvient malheureusement pas à fournir un comportement universellement acceptable). Merci
manger
3

La méthode np.argpartitionne renvoie que les k plus grands indices, effectue un tri local et est plus rapide que np.argsort(effectuer un tri complet) lorsque le tableau est assez grand. Mais les indices retournés ne sont PAS dans l'ordre croissant / décroissant . Disons avec un exemple:

Entrez la description de l'image ici

Nous pouvons voir que si vous voulez un ordre k croissant, les meilleurs index np.argpartitionne retourneront pas ce que vous voulez.

En plus de faire un tri manuellement après np.argpartition, ma solution est d'utiliser PyTorch, torch.topkun outil pour la construction de réseaux neuronaux, fournissant des API de type NumPy avec prise en charge CPU et GPU. Il est aussi rapide que NumPy avec MKL et offre un boost GPU si vous avez besoin de grands calculs matriciels / vectoriels.

Le code strict des k premiers indices ascendants / descendants sera:

Entrez la description de l'image ici

Notez qu'il torch.topkaccepte un tenseur de torche et renvoie à la fois les k premières valeurs et les k premiers indices en type torch.Tensor. Similaire à np, torch.topk accepte également un argument d'axe afin que vous puissiez gérer des tableaux / tenseurs multidimensionnels.

futuriste
la source
2

Utilisation:

from operator import itemgetter
from heapq import nlargest
result = nlargest(N, enumerate(your_list), itemgetter(1))

Maintenant, la resultliste contiendrait N tuples ( index, value) où valueest maximisé.

off99555
la source
2

Utilisation:

def max_indices(arr, k):
    '''
    Returns the indices of the k first largest elements of arr
    (in descending order in values)
    '''
    assert k <= arr.size, 'k should be smaller or equal to the array size'
    arr_ = arr.astype(float)  # make a copy of arr
    max_idxs = []
    for _ in range(k):
        max_element = np.max(arr_)
        if np.isinf(max_element):
            break
        else:
            idx = np.where(arr_ == max_element)
        max_idxs.append(idx)
        arr_[idx] = -np.inf
    return max_idxs

Il fonctionne également avec des tableaux 2D. Par exemple,

In [0]: A = np.array([[ 0.51845014,  0.72528114],
                     [ 0.88421561,  0.18798661],
                     [ 0.89832036,  0.19448609],
                     [ 0.89832036,  0.19448609]])
In [1]: max_indices(A, 8)
Out[1]:
    [(array([2, 3], dtype=int64), array([0, 0], dtype=int64)),
     (array([1], dtype=int64), array([0], dtype=int64)),
     (array([0], dtype=int64), array([1], dtype=int64)),
     (array([0], dtype=int64), array([0], dtype=int64)),
     (array([2, 3], dtype=int64), array([1, 1], dtype=int64)),
     (array([1], dtype=int64), array([1], dtype=int64))]

In [2]: A[max_indices(A, 8)[0]][0]
Out[2]: array([ 0.89832036])
X Æ A-12
la source
Fonctionne bien, mais donne plus de résultats si vous avez des valeurs en double (maximum) dans votre tableau A. Je m'attendrais exactement à k résultats mais en cas de valeurs en double, vous obtenez plus de k résultats.
Guido
J'ai légèrement modifié le code. La liste des indices renvoyée a une longueur exactement égale à k. Si vous avez des doublons, ils sont regroupés en un seul tuple.
X Æ A-12
1

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

Je ne sais rien de ce module; Je viens de googler numpy partial sort.

Katriel
la source
Je ne trouve aucune fonction de tri partiel dans le goulot d'étranglement, il y a une fonction de partition, mais cela ne trie pas
nbecker
1

Ce qui suit est un moyen très simple de voir le maximum d'éléments et ses positions. Voici axisle domaine; axis= 0 signifie le nombre maximum par colonne et axis= 1 signifie le nombre maximum par ligne pour le cas 2D. Et pour des dimensions supérieures, cela dépend de vous.

M = np.random.random((3, 4))
print(M)
print(M.max(axis=1), M.argmax(axis=1))
libéral
la source
J'ai utilisé ce lien jakevdp.github.io/PythonDataScienceHandbook/…
libéral
0

Je l'ai trouvé le plus intuitif à utiliser np.unique.

L'idée est que la méthode unique renvoie les indices des valeurs d'entrée. Ensuite, à partir de la valeur unique maximale et des indices, la position des valeurs d'origine peut être recréée.

multi_max = [1,1,2,2,4,0,0,4]
uniques, idx = np.unique(multi_max, return_inverse=True)
print np.squeeze(np.argwhere(idx == np.argmax(uniques)))
>> [4 7]
phi
la source
0

Je pense que le moyen le plus efficace en termes de temps est d'itérer manuellement le tableau et de conserver un min-tas de taille k, comme d'autres l'ont mentionné.

Et je propose également une approche par force brute:

top_k_index_list = [ ]
for i in range(k):
    top_k_index_list.append(np.argmax(my_array))
    my_array[top_k_index_list[-1]] = -float('inf')

Définissez le plus grand élément sur une grande valeur négative après avoir utilisé argmax pour obtenir son index. Et puis le prochain appel d'argmax retournera le deuxième plus grand élément. Et vous pouvez enregistrer la valeur d'origine de ces éléments et les récupérer si vous le souhaitez.

Zhenghao Zhao
la source
0

Ce code fonctionne pour un tableau matriciel numpy:

mat = np.array([[1, 3], [2, 5]]) # numpy matrix

n = 2  # n
n_largest_mat = np.sort(mat, axis=None)[-n:] # n_largest 
tf_n_largest = np.zeros((2,2), dtype=bool) # all false matrix
for x in n_largest_mat: 
  tf_n_largest = (tf_n_largest) | (mat == x) # true-false  

n_largest_elems = mat[tf_n_largest] # true-false indexing 

Cela produit une indexation matricielle n_largest vraie-fausse qui fonctionne également pour extraire les éléments n_largest d'un tableau matriciel

Yi Xiang Chong
la source