Come applicare l'output di numpy.argpartition per gli array 2-D?

Question

Ho un array numerico 2d più grande, e voglio estrarre i 10 elementi più bassi di ogni riga, nonché i loro indici. Dal momento che il mio array è più ampio, preferirei non ordinare l'intero array.

ho sentito sulla funzione argpartition(), con la quale posso ottenere gli indici dei più bassi 10 elementi:

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

Nota che argpartition() partizioni asse -1 di default, che è quello che voglio. Il risultato qui ha la stessa forma di myBigArray che contiene gli indici nelle rispettive file in modo tale che i primi 10 indici puntino ai 10 valori più bassi.

Come è possibile estrarre gli elementi di myBigArray corrispondenti a tali indici?

L'evidente indicizzazione di fantasia come myBigArray[top10indexes] o myBigArray[:,top10indexes] fa qualcosa di molto diverso. Potrei anche usare list comprehension, qualcosa di simile:

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

ma sarebbe incorrere in una performance ha colpito l'iterazione righe numpy e convertendo il risultato di nuovo a un array.

nb: Potrei semplicemente usare np.partition() per ottenere i valori, e potrebbero anche corrispondere agli indici (o potrebbero non ..), ma non voglio fare la partizione due volte se posso evitarlo.

Answer 1

È possibile evitare di utilizzare le copie appiattito e la necessità di estrarre tutti i valori facendo:

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

Per NumPy> = 1.9.0 questo sarà molto efficace e paragonabile a np.take().