¿Cómo obtengo índices de N valores máximos en una matriz NumPy?

NumPy propone una forma de obtener el índice del valor máximo de una matriz a través de np.argmax.

Me gustaría algo similar, pero devolviendo los índices de los Nvalores máximos.

Por ejemplo, si tengo una matriz, [1, 3, 2, 4, 5], function(array, n=3)volvería los índices [4, 3, 1]que corresponden a los elementos [5, 4, 3].

Lo más simple que he podido encontrar es:

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

Esto implica un tipo completo de la matriz. Me pregunto si numpyproporciona una forma integrada de hacer una clasificación parcial; Hasta ahora no he podido encontrar uno.

Si esta solución resulta ser demasiado lenta (especialmente para los pequeños n), puede valer la pena buscar codificar algo en Cython .

Las versiones más recientes de NumPy (1.8 y superiores) tienen una función llamada argpartitionpara esto. Para obtener los índices de los cuatro elementos más grandes, haga

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

A diferencia argsort, esta función se ejecuta en tiempo lineal en el peor de los casos, pero los índices devueltos no están ordenados, como se puede ver en el resultado de la evaluación a[ind]. Si también lo necesitas, clasifícalos después:

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

Para obtener los elementos top- k en orden ordenado de esta manera se necesita tiempo O ( n + k log k ).

Más simple aún:

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

donde n es el número de valores máximos.

Utilizar:

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

Para listas regulares de Python:

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

Si usa Python 2, use en xrangelugar de range.

Fuente: heapq - Algoritmo de cola de montón

Si está trabajando con una matriz multidimensional, deberá aplanar y desentrañar los índices:

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)

Por ejemplo:

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

Si no le importa el orden de los elementos K-th más grandes que puede usar argpartition, que deberían funcionar mejor que una clasificación completa 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])

Los créditos van a esta pregunta .

Ejecuté algunas pruebas y parece que tiene un argpartitionrendimiento superior a argsortmedida que aumenta el tamaño de la matriz y el valor de K.

Para las matrices multidimensionales, puede usar la axispalabra clave para aplicar la partición a lo largo del eje esperado.

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

Y para agarrar los artículos:

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

Pero tenga en cuenta que esto no devolverá un resultado ordenado. En ese caso, puede usar a lo np.argsort()largo del eje 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)

Aquí hay un ejemplo:

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

Esto será más rápido que una clasificación completa, dependiendo del tamaño de su matriz original y el tamaño de su selección:

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

Por supuesto, implica la manipulación de su matriz original. Lo que puede solucionar (si es necesario) haciendo una copia o reemplazando los valores originales. ... lo que sea más barato para su caso de uso.

El método np.argpartitionsolo devuelve los k índices más grandes, realiza una ordenación local y es más rápido que np.argsort(realizar una ordenación completa) cuando la matriz es bastante grande. Pero los índices devueltos NO están en orden ascendente / descendente . Digamos con un ejemplo:

Podemos ver que si desea un orden ascendente estricto de k índices superiores, np.argpartitionno le devolverá lo que desea.

Además de hacer una clasificación manual después de np.argpartition, mi solución es usar PyTorch, torch.topkuna herramienta para la construcción de redes neuronales, que proporciona API similares a NumPy con soporte para CPU y GPU. Es tan rápido como NumPy con MKL, y ofrece un impulso de GPU si necesita grandes cálculos de matriz / vector.

El estricto código de los principales índices de ascenso / descenso será:

Tenga en cuenta que torch.topkacepta un tensor de antorcha y devuelve los valores k superiores y los índices k superiores en tipo torch.Tensor. Similar a np, torch.topk también acepta un argumento de eje para que pueda manejar matrices / tensores multidimensionales.

Utilizar:

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

Ahora la resultlista contendría N tuplas ( index, value) donde valueestá maximizada.

Utilizar:

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

También funciona con matrices 2D. Por ejemplo,

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 tiene una función de ordenación parcial, si el gasto de ordenar la matriz completa solo para obtener los N valores más grandes es demasiado grande.

_{No sé nada sobre este módulo; Yo solo busqué en Google numpy partial sort.}

La siguiente es una manera muy fácil de ver los elementos máximos y sus posiciones. Aquí axisestá el dominio; axis= 0 significa número máximo de columna y axis= 1 significa número máximo de fila para el caso 2D. Y para dimensiones superiores depende de ti.

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

Lo encontré más intuitivo de usar np.unique.

La idea es que el método único devuelve los índices de los valores de entrada. Luego, a partir del valor único máximo y las indicaciones, se puede recrear la posición de los valores originales.

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]

Creo que la forma más eficiente de tiempo es iterar manualmente a través de la matriz y mantener un montón mínimo de tamaño k, como han mencionado otras personas.

Y también se me ocurre un enfoque de fuerza 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')

Establezca el elemento más grande en un valor negativo grande después de usar argmax para obtener su índice. Y luego la próxima llamada de argmax devolverá el segundo elemento más grande. Y puede registrar el valor original de estos elementos y recuperarlos si lo desea.

Este código funciona para una matriz de matriz 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 

Esto produce una indexación de matriz n_largest verdadero-falso que también funciona para extraer elementos n_largest de una matriz de matriz