Ottieni max o min elementi n dall'array numpy? (Preferibilmente non appiattito)

Question

So che posso ottenere min o valori massimi con:

max(matrix) 
min(matrix) 

su un NumPy matrice/vettore. Gli indici per tali valori vengono restituiti da:

argmax(matrix) 
argmin(matrix) 

Così ad es. quando ho una matrice 5x5:

a = np.arange(5*5).reshape(5, 5) + 10 

# array([[10, 11, 12, 13, 14], 
#  [15, 16, 17, 18, 19], 
#  [20, 21, 22, 23, 24], 
#  [25, 26, 27, 28, 29], 
#  [30, 31, 32, 33, 34]]) 

ho potuto ottenere il valore massimo attraverso:

In [86]: np.max(a) # getting the max-value out of a 
Out[86]: 34 

In [87]: np.argmax(a) # index of max-value 34 is 24 if array a were flattened 
Out[87]: 24 

... ma qual è il modo più efficace per ottenere il massimo o il minimo n elementi?

Quindi diciamo su a Voglio avere i 5 elementi più in alto e 5 in basso. Questo dovrebbe restituirmi [30, 31, 32, 33, 34] per i 5 valori più alti rispettivamente [20, 21, 22, 23, 24] per i loro indici. Allo stesso modo [10, 11, 12, 13, 14] per i 5 valori più bassi e [0, 1, 2, 3, 4] per gli indici dei 5 elementi più bassi.

Quale sarebbe una soluzione efficiente e ragionevole per questo?

La mia prima idea eraappiattimento e l'ordinamento l'array e prendendo gli ultimi e primi 5 valori. Successivamente, cerco la matrice 2D originale per gli indici di tali valori. Anche se questa procedura funziona appiattendo + l'ordinamento non è molto efficiente ... qualcuno conosce una soluzione più veloce?

Inoltre mi piacerebbe avere gli indici dell'array 2D originale e non quello di appiattimento. Quindi invece di 24 restituito da np.argmax(a) mi piacerebbe avere (4, 4).

Answer 1

Il modo standard per ottenere gli indici dei valori più grandi o più piccoli in un array è utilizzare np.argpartition. Questa funzione utilizza un algoritmo introselect e gira con complessità lineare - questo funziona meglio dell'ordinamento completo per array più grandi (che è tipicamente O (n log n)).

Per impostazione predefinita questa funzione funziona lungo l'ultimo asse dell'array. Per considerare un intero array, è necessario utilizzare ravel(). Ad esempio, ecco una matrice casuale a:

>>> a = np.random.randint(0, 100, size=(5, 5)) 
>>> a 
array([[60, 68, 86, 66, 9], 
     [66, 26, 83, 87, 50], 
     [41, 26, 0, 55, 9], 
     [57, 80, 71, 50, 22], 
     [94, 30, 95, 99, 76]]) 

Poi per ottenere gli indici dei cinque valori più grandi del (appiattito) matrice 2D, uso:

>>> i = np.argpartition(a.ravel(), -5)[-5:] # argpartition(a.ravel(), 5)[:5] for smallest 
>>> i 
array([ 2, 8, 22, 23, 20]) 

tornare indici 2D corrispondenti di queste posizioni nella a, utilizzare unravel_index:

>>> i2d = np.unravel_index(i, a.shape) 
>>> i2d 
(array([0, 1, 4, 4, 4]), array([2, 3, 2, 3, 0])) 

quindi indicizzare a con i2d riconsegna i cinque maggiori valori:

>>> a[i2d] 
array([86, 87, 95, 99, 94])