Wie erhalte ich Indizes von N maximalen Werten in einem NumPy-Array?

Question

NumPy bietet eine Möglichkeit, den Index des maximalen Wertes eines Arrays über np.argmax .

Ich würde gerne etwas Ähnliches machen, aber die Indizes der N Höchstwerte.

Zum Beispiel, wenn ich ein Array habe, [1, 3, 2, 4, 5] entonces nargmax(array, n=3) würde die Indizes zurückgeben [4, 3, 1] die den folgenden Elementen entsprechen [5, 4, 3] .

Answer 1

Drei Antworten im Vergleich für einfaches und schnelles Coding

Geschwindigkeit war für meine Bedürfnisse wichtig, daher habe ich drei Antworten auf diese Frage getestet.

Der Code aus diesen drei Antworten wurde nach Bedarf für meinen speziellen Fall geändert.

Dann habe ich die Geschwindigkeit der einzelnen Methoden verglichen.

Kodierungstechnisch:

1. Die Antwort von NPE war die nächst elegantere und für meine Bedürfnisse ausreichend schnell.
2. Die Antwort von Fred Foos erforderte die meisten Anpassungen für meine Bedürfnisse, war aber auch die schnellste. Ich habe mich für diese Antwort entschieden, weil sie zwar mehr Arbeit erforderte, aber nicht allzu schlecht war und erhebliche Geschwindigkeitsvorteile bot.
3. Die Antwort von off99555 war die eleganteste, aber auch die langsamste.

Vollständiger Code für Test und Vergleiche

import numpy as np
import time
import random
import sys
from operator import itemgetter
from heapq import nlargest

''' Fake Data Setup '''
a1 = list(range(1000000))
random.shuffle(a1)
a1 = np.array(a1)

''' ################################################ '''
''' NPE's Answer Modified A Bit For My Case '''
t0 = time.time()
indices = np.flip(np.argsort(a1))[:5]
results = []
for index in indices:
    results.append((index, a1[index]))
t1 = time.time()
print("NPE's Answer:")
print(results)
print(t1 - t0)
print()

''' Fred Foos Answer Modified A Bit For My Case'''
t0 = time.time()
indices = np.argpartition(a1, -6)[-5:]
results = []
for index in indices:
    results.append((a1[index], index))
results.sort(reverse=True)
results = [(b, a) for a, b in results]
t1 = time.time()
print("Fred Foo's Answer:")
print(results)
print(t1 - t0)
print()

''' off99555's Answer - No Modification Needed For My Needs '''
t0 = time.time()
result = nlargest(5, enumerate(a1), itemgetter(1))
t1 = time.time()
print("off99555's Answer:")
print(result)
print(t1 - t0)

Ausgabe mit Speed Reports

Die Antwort der NPE:

[(631934, 999999), (788104, 999998), (413003, 999997), (536514, 999996), (81029, 999995)]
0.1349949836730957

Die Antwort von Fred Foo:

[(631934, 999999), (788104, 999998), (413003, 999997), (536514, 999996), (81029, 999995)]
0.011161565780639648

off99555's Antwort:

[(631934, 999999), (788104, 999998), (413003, 999997), (536514, 999996), (81029, 999995)]
0.439760684967041

Answer 2

Wenn Sie sich nicht für die Bestellung der K-ten größten Elemente können Sie verwenden argpartition was besser sein sollte als ein vollständiges Durchsortieren 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])

Die Credits gehen an diese Frage .

Ich habe ein paar Tests durchgeführt und es sieht so aus argpartition besser abschneidet als argsort wenn die Größe des Feldes und der Wert von K zunehmen.

Answer 3

Für mehrdimensionale Arrays können Sie die axis Schlüsselwort, um die Partitionierung entlang der erwarteten Achse anzuwenden.

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

Und für das Ergreifen der Gegenstände:

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

Beachten Sie jedoch, dass dies kein sortiertes Ergebnis liefert. In diesem Fall können Sie np.argsort() entlang der vorgesehenen Achse:

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

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

Hier ist ein Beispiel:

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]])

Answer 4

Methode np.argpartition gibt nur die k größten Indizes zurück, führt eine lokale Sortierung durch und ist schneller als np.argsort (Durchführung einer vollständigen Sortierung), wenn das Array recht groß ist. Aber die zurückgegebenen Indizes sind NICHT in aufsteigender/absteigender Reihenfolge . Nehmen wir ein Beispiel:

Wir können sehen, dass, wenn Sie eine streng aufsteigende Reihenfolge top k Indizes wollen, np.argpartition wird nicht das zurückgeben, was Sie wollen.

Abgesehen von einer manuellen Sortierung nach np.argpartition, ist meine Lösung, PyTorch zu verwenden, torch.topk ein Werkzeug zur Konstruktion neuronaler Netze, das NumPy-ähnliche APIs mit CPU- und GPU-Unterstützung bietet. Es ist genauso schnell wie NumPy mit MKL und bietet einen GPU-Boost, wenn Sie große Matrix-/Vektorberechnungen benötigen.

Streng aufsteigende/absteigende Top-K-Indizes werden codiert:

Beachten Sie, dass torch.topk akzeptiert einen Torch-Tensor und gibt sowohl Top-k-Werte als auch Top-k-Indizes vom Typ torch.Tensor . Ähnlich wie bei np akzeptiert torch.topk auch ein Achsenargument, so dass Sie mit mehrdimensionalen Arrays/Tensoren arbeiten können.

Answer 5

Dies ist schneller als eine vollständige Sortierung, abhängig von der Größe des ursprünglichen Arrays und der Größe der Auswahl:

>>> 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])

Das bedeutet natürlich, dass Sie Ihre ursprüngliche Anordnung manipulieren müssen. Das können Sie (falls nötig) durch Anfertigen einer Kopie oder Ersetzen der Originalwerte beheben. ...je nachdem, was für Ihren Anwendungsfall billiger ist.