¿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 N valores máximos.

Por ejemplo, si tengo una matriz, [1, 3, 2, 4, 5] , function(array, n=3) devolverí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 numpy proporciona una forma integrada de hacer una ordenación parcial; Hasta ahora no he podido encontrar uno.

Si esta solución resulta demasiado lenta (especialmente para las pequeñas n ), puede valer la pena ver cómo codificar algo en Cython .

Las versiones más nuevas de NumPy (1.8 y superiores) tienen una función llamada argpartition para 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 de argsort , esta función se ejecuta en tiempo lineal en el peor de los casos, pero los índices devueltos no están ordenados, como puede verse en el resultado de evaluar a[ind] . Si también necesitas eso, ordénalos después:

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

Para obtener los mejores elementos de k en orden, de esta forma, se tarda 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 las listas regulares de Python:

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

Si usas Python 2, usa xrange lugar de range .

Fuente: heapq – Algoritmo de cola de stack

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 te importa el orden de los elementos K-th más grandes, puedes usar argpartition , que debería funcionar mejor que una ordenación completa a través de 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 .

argpartition algunas pruebas y parece que argpartition supera a argsort medida que argsort el tamaño de la matriz y el valor de K.

Para matrices multidimensionales, puede utilizar la palabra clave axis 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 objetos:

 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 np.argsort() largo del eje deseado:

 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 ordenació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. Que podría arreglar (si es necesario) haciendo una copia o reemplazando los valores originales. … lo que sea más barato para su caso de uso.

bottleneck tiene una función de clasificació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 de este módulo; Acabo de googled numpy partial sort .

Utilizar:

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

Ahora la lista de result contendría N tuples ( index , value ) donde el value se maximiza.

El método np.argpartition solo devuelve los k índices más grandes, realiza una ordenación local y es más rápido que np.argsort (realiza 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:

Introduzca la descripción de la imagen aquí

Podemos ver que si desea un índice ascendente k de orden ascendente estricto, np.argpartition no devolverá lo que desea.

Además de hacer una ordenación manual después de np.argpartition, mi solución es usar PyTorch, torch.topk , una herramienta para la construcción de redes neuronales, que proporciona API de tipo 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 matrices / vectores.

El código de índices de ascenso / descenso superior k será:

Introduzca la descripción de la imagen aquí

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

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

La siguiente es una manera muy fácil de ver los elementos máximos y sus posiciones. Aquí el axis es el dominio; axis = 0 significa el número máximo sabio de la columna y axis = 1 significa el número máximo sabio de la 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)) 

Me pareció más intuitivo usar np.unique .

La idea es que el método único devuelva 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 min-heap 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 siguiente llamada de argmax devolverá el segundo elemento más grande. Y puede registrar el valor original de estos elementos y recuperarlos si lo desea.