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

485

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

Alexis Métaireau
fuente
44
Su pregunta no está realmente bien definida. Por ejemplo, ¿cuáles serían los índices (que esperas) para ser array([5, 1, 5, 5, 2, 3, 2, 4, 1, 5]), pizca n= 3? ¿Cuál de todas las alternativas, como [0, 2, 3], [0, 2, 9], ...sería la correcta? Por favor, explique más sobre sus requisitos específicos. Gracias
comer
@eat, realmente no me importa cuál se supone que debe devolverse en este caso específico. Incluso si parece lógico devolver el primero encontrado, no es un requisito para mí.
Alexis Métaireau
argsortpodría ser una alternativa viable si no le importa el orden de las devoluciones devueltas. Vea mi respuesta a continuación.
azul

Respuestas:

349

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 .

NPE
fuente
1
¿Podría la línea 3 escribirse de manera equivalente como arr.argsort()[-1:-4:-1]? Lo probé en el intérprete y aparece el mismo resultado, pero me pregunto si no está roto por algún ejemplo.
abroekhof
44
@abroekhof Sí, eso debería ser equivalente para cualquier lista o matriz. Alternativamente, esto podría hacerse sin la inversión mediante el uso np.argsort(-arr)[:3], que me parece más legible y al grano.
askewchan
66
¿Qué significa [:: - 1]? @NPE
1a1a11a
@ 1a1a11a significa revertir una matriz (literalmente, toma una copia de una matriz de min sin restricciones a max sin restricciones en un orden inverso)
FizBack
15
arr.argsort()[::-1][:n]es mejor porque regresa vacío para en n=0lugar de la matriz completa
abora
600

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

Fred Foo
fuente
27
@varela se argpartitionejecuta en tiempo lineal, O (n), utilizando el algoritmo de introselección . La clasificación posterior solo maneja k elementos, de modo que se ejecuta en O (k log k).
Fred Foo
2
Si alguien se pregunta cómo funciona exactamente np.argpartitiony su algoritmo hermano, np.partitionhay una explicación más detallada en la pregunta vinculada: stackoverflow.com/questions/10337533/…
Ramon Martinez
77
@FredFoo: ¿por qué usaste -4? ¿Hiciste eso para comenzar hacia atrás? (¡ya que k ser positivo o negativo funciona igual para mí! ¡Solo imprime los números más pequeños primero!
Rika
2
Uso @LKT a=np.array([9, 4, 4, 3, 3, 9, 0, 4, 6, 0])porque las listas normal de Python no son compatibles con la indexación de las listas, a diferencia denp.array
Marawan Okasha
2
@Umangsinghal np.argpartitiontoma un axisargumento opcional . Para encontrar los índices de los valores n superiores para cada fila:np.argpartition(a, -n, axis=1)[-n:]
Ralph
48

Más simple aún:

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

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

Ketan
fuente
77
¿Se puede hacer esto para una matriz 2d? Si no, ¿quizás sabes cómo?
Andrew Hundt el
2
@AndrewHundt: simplemente use (-arr) .argsort (axis = -1) [:,: n]
MiniQuark
2
similar sería arr[arr.argsort()[-n:]]en lugar de negar la matriz, acaba de tomar una rebanada de los últimos elementos n
loganjones16
35

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

anishpatel
fuente
2
No hay necesidad de un bucle en absoluto aquí: heapq.nlargest(3, xrange(len(a)), a.take). Para las listas de Python podemos usar en .__getitem__lugar de .take.
Ashwini Chaudhary
Para las matrices n-dimensionales Aen general: heapq.nlargest(3, range(len(A.ravel())), A.ravel().take). (Espero que esto solo funcione en vistas, ver también ( ravel vs flatten] ( stackoverflow.com/a/28930580/603003 )).
ComFreek
31

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])
danvk
fuente
9

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.

azul
fuente
7

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]])
Kasramvd
fuente
Creo que puede simplificar la indexación aquí mediante el uso np.take_along_axis(que probablemente no existía cuando respondió esta pregunta)
Eric
4

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.

Pablo
fuente
FWIW, su solución no proporcionará una solución inequívoca en todas las situaciones. OP debería describir cómo manejar estos casos inequívocos. Gracias
comer
@eat La pregunta del OP es un poco ambigua. Sin embargo, una implementación no está realmente abierta a interpretación. :) El OP simplemente debe referirse a la definición de np.argmax docs.scipy.org/doc/numpy/reference/generated/numpy.argmax.html para asegurarse de que esta solución específica cumpla con los requisitos. Es posible que cualquier solución que cumpla con los requisitos establecidos por el OP sea aceptable ..
Paul
Bueno, uno podría considerar que la implementación de argmax(.)no es ambigua también. (En mi humilde opinión, intenta seguir algún tipo de lógica de cortocircuito, pero desafortunadamente no proporciona un comportamiento universalmente aceptable). Gracias
comer
3

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:

Ingrese la descripción de la imagen aquí

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á:

Ingrese la descripción de la imagen aquí

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.

futuro
fuente
2

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.

off99555
fuente
2

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])
X Æ A-12
fuente
Funciona bien, pero da más resultados si tiene valores duplicados (máximos) en su matriz A. Esperaría exactamente k resultados, pero en caso de valores duplicados, obtendrá más de k resultados.
Guido
Modifiqué ligeramente el código. La lista de índices que se devuelve tiene una longitud igual exactamente a k. Si tiene duplicados, se agrupan en una sola tupla.
X Æ A-12
1

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.

Katriel
fuente
No encuentro ninguna función de clasificación parcial en el cuello de botella, hay una función de partición, pero esto no significa especie
nbecker
1

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))
liberal
fuente
Usé
liberal
0

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]
fi
fuente
0

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.

Zhenghao Zhao
fuente
0

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

Yi Xiang Chong
fuente