Wie wendet man die Ausgabe von numpy.argpartition für 2D-Arrays an?

Question

Ich habe eine große 2d numpy Array, und ich möchte die niedrigsten 10 Elemente jeder Zeile sowie ihre Indizes extrahieren. Da mein Array größer ist, würde ich es vorziehen, das gesamte Array nicht zu sortieren.

ich über die argpartition() Funktion gehört, mit dem ich die Indizes der untersten 10 Elemente erhalten können:

top10indexes = np.argpartition(myBigArray,10)[:,:10] 

Beachten Sie, dass argpartition() Partitionen Achse -1 standardmäßig, das ist das, was ich will. Das Ergebnis hat dieselbe Form wie myBigArray und enthält Indizes in die entsprechenden Zeilen, so dass die ersten 10 Indizes auf die 10 niedrigsten Werte zeigen.

Wie kann ich jetzt die Elemente myBigArray entsprechend diesen Indizes extrahieren?

Offensichtliche Phantasie Indizierung wie myBigArray[top10indexes] oder myBigArray[:,top10indexes] tun etwas ganz anderes. Ich könnte auch Listenkomprehensionen, so etwas wie verwenden:

array([row[idxs] for row,idxs in zip(myBigArray,top10indexes)]) 

aber das wäre eine Leistung entsteht numpy Reihen getroffen laufen und das Ergebnis zurück auf ein Array zu konvertieren.

nb: Ich könnte einfach np.partition() verwenden, um die Werte zu erhalten, und sie können sogar den Indizes entsprechen (oder nicht ..), aber ich möchte die Partition nicht zweimal machen, wenn ich es vermeiden kann.

Answer 1

Sie können mit den abgeflachten Kopien und die Notwendigkeit vermeiden, alle Werte zu extrahieren, indem Sie:

num = 10 
top = np.argpartition(myBigArray, num, axis=1)[:, :num] 
myBigArray[np.arange(myBigArray.shape[0])[:, None], top] 

Für NumPy> = 1.9.0 dies wird sehr effizient und vergleichbar mit np.take().