Come posso ottenere indici di N valori massimi in un array NumPy?

NumPy propone un modo per ottenere l'indice del valore massimo di un array tramite np.argmax.

Vorrei una cosa simile, ma restituendo gli indici dei Nvalori massimi.

Ad esempio, se ho un array, [1, 3, 2, 4, 5], function(array, n=3)restituirebbe gli indici [4, 3, 1]che corrispondono agli elementi [5, 4, 3].

Il più semplice che sono stato in grado di inventare è:

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

Ciò comporta una sorta di array completo. Mi chiedo se numpyfornisce un modo integrato per fare un ordinamento parziale; finora non sono stato in grado di trovarne uno.

Se questa soluzione risulta troppo lenta (specialmente per i piccoli n), può valere la pena di cercare di scrivere qualcosa su Cython .

Le versioni più recenti di NumPy (1.8 e successive) hanno una funzione chiamata argpartitionper questo. Per ottenere gli indici dei quattro elementi più grandi, fallo

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

Diversamente argsort, questa funzione viene eseguita nel tempo lineare nel peggiore dei casi, ma gli indici restituiti non vengono ordinati, come si può vedere dal risultato della valutazione a[ind]. Se ne hai bisogno anche tu, ordinali in seguito:

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

Per ottenere i top- k elementi in modo ordinato in questo modo prende O ( n + k log k tempo).

Ancora più semplice:

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

dove n è il numero di valori massimi.

Uso:

>>> import heapq
>>> import numpy
>>> a = numpy.array([1, 3, 2, 4, 5])
>>> heapq.nlargest(3, range(len(a)), a.take)
[4, 3, 1]

Per gli elenchi Python regolari:

>>> a = [1, 3, 2, 4, 5]
>>> heapq.nlargest(3, range(len(a)), a.__getitem__)
[4, 3, 1]

Se usi Python 2, usa xrangeinvece di range.

Fonte: heapq - algoritmo della coda di heap

Se ti capita di lavorare con un array multidimensionale, dovrai appiattire e svelare gli indici:

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)

Per esempio:

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

Se non ti interessa l' ordine degli elementi più grandi del K-esimo che puoi usare argpartition, che dovrebbe funzionare meglio di un ordinamento completo 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])

I crediti vanno a questa domanda .

Ho eseguito alcuni test e sembra argpartitionsovraperformare argsortall'aumentare della dimensione dell'array e del valore di K.

Per gli array multidimensionali è possibile utilizzare la axisparola chiave per applicare il partizionamento lungo l'asse previsto.

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

E per afferrare gli oggetti:

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

Si noti che ciò non restituirà un risultato ordinato. In tal caso è possibile utilizzare np.argsort()lungo l'asse previsto:

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

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

Ecco un esempio:

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

Questo sarà più veloce di un ordinamento completo a seconda delle dimensioni dell'array originale e delle dimensioni della selezione:

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

Naturalmente, comporta la manomissione dell'array originale. Che è possibile correggere (se necessario) eseguendo una copia o sostituendo i valori originali. ... qualunque sia il più economico per il tuo caso d'uso.

Il metodo np.argpartitionrestituisce solo i k indici più grandi, esegue un ordinamento locale ed è più veloce di np.argsort(eseguendo un ordinamento completo) quando l'array è piuttosto grande. Ma gli indici restituiti NON sono in ordine crescente / decrescente . Diciamo con un esempio:

Possiamo vedere che se si desidera un ordine ascendente rigoroso in cima agli indici k, np.argpartitionnon restituirà ciò che si desidera.

Oltre a fare una sorta di ordinamento manualmente dopo np.argpartition, la mia soluzione è utilizzare PyTorch, torch.topkuno strumento per la costruzione di reti neurali, che fornisce API simili a NumPy con supporto sia CPU che GPU. È veloce come NumPy con MKL e offre un aumento della GPU se hai bisogno di calcoli con matrice / vettore di grandi dimensioni.

Il codice degli indici top k ascendenti / discendenti rigorosi sarà:

Si noti che torch.topkaccetta un tensore della torcia e restituisce sia i valori k principali che gli indici k superiori nel tipo torch.Tensor. Simile a np, torch.topk accetta anche un argomento dell'asse in modo da poter gestire array / tensori multidimensionali.

Uso:

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

Ora l' resultelenco conterrebbe N tuple ( index, value) dove valueè ingrandito.

Uso:

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

Funziona anche con array 2D. Per esempio,

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

bottleneck ha una funzione di ordinamento parziale, se la spesa per ordinare l'intero array solo per ottenere i valori N più grandi è troppo grande.

_{Non so nulla di questo modulo; Ho appena cercato su Google numpy partial sort.}

Di seguito è riportato un modo molto semplice per visualizzare gli elementi massimi e le relative posizioni. Ecco axisil dominio; axis= 0 indica il numero massimo per colonna e axis= 1 indica il numero massimo per riga per il caso 2D. E per dimensioni superiori dipende da te.

M = np.random.random((3, 4))
print(M)
print(M.max(axis=1), M.argmax(axis=1))

L'ho trovato molto intuitivo da usare np.unique.

L'idea è che il metodo univoco restituisca gli indici dei valori di input. Quindi, dal valore univoco massimo e dalle indicazioni, è possibile ricreare la posizione dei valori originali.

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]

Penso che il modo più efficiente in termini di tempo sia iterare manualmente l'array e mantenere un heap minimo di dimensioni k, come altri hanno già detto.

E ho anche escogitato un approccio a forza bruta:

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

Impostare l'elemento più grande su un valore negativo elevato dopo aver usato argmax per ottenere il suo indice. E quindi il prossimo richiamo di argmax restituirà il secondo elemento più grande. E puoi registrare il valore originale di questi elementi e recuperarli se vuoi.

Questo codice funziona per una matrice di matrice 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 

Questo produce un indice di matrice n_largest vero-falso che funziona anche per estrarre elementi n_largest da un array di matrici