Obtenga N valores e índices máximos a lo largo de un eje en una matriz NumPy

Creo que esta es una pregunta fácil para usuarios experimentados.

Tengo una matriz de puntuación. El índice bruto corresponde a las muestras y el índice de columna corresponde a los elementos. Por ejemplo,

score_matrix = [[ 1. , 0.3, 0.4], [ 0.2, 0.6, 0.8], [ 0.1, 0.3, 0.5]] 

Quiero obtener los mejores índices M de ítems para cada muestra. También quiero conseguir los mejores puntajes. Por ejemplo,

 top2_ind = [[0, 2], [2, 1], [2, 1]] top2_score = [[1. , 0.4], [0,8, 0.6], [0.5, 0.3]] 

¿Cuál es la mejor manera de hacerlo usando numpy?

Yo usaría argsort() :

 top2_ind = score_matrix.argsort()[:,::-1][:,:2] 

Es decir, producir una matriz que contenga los índices que clasificarían score_matrix :

 array([[1, 2, 0], [0, 1, 2], [0, 1, 2]]) 

Luego invierta las columnas con ::-1 , luego tome las dos primeras columnas con :2 :

 array([[0, 2], [2, 1], [2, 1]]) 

Luego similar pero con np.sort() normal para obtener los valores:

 top2_score = np.sort(score_matrix)[:,::-1][:,:2] 

Que siguiendo la misma mecánica que el anterior, te da:

 array([[ 1. , 0.4], [ 0.8, 0.6], [ 0.5, 0.3]]) 

Aquí hay un enfoque usando np.argpartition

 idx = np.argpartition(a,range(M))[:,:-M-1:-1] # topM_ind out = a[np.arange(a.shape[0])[:,None],idx] # topM_score 

Ejecución de la muestra

 In [343]: a Out[343]: array([[ 1. , 0.3, 0.4], [ 0.2, 0.6, 0.8], [ 0.1, 0.3, 0.5]]) In [344]: M = 2 In [345]: idx = np.argpartition(a,range(M))[:,:-M-1:-1] In [346]: idx Out[346]: array([[0, 2], [2, 1], [2, 1]]) In [347]: a[np.arange(a.shape[0])[:,None],idx] Out[347]: array([[ 1. , 0.4], [ 0.8, 0.6], [ 0.5, 0.3]]) 

Alternativamente, posiblemente más lento, pero un código un poco más corto para obtener idx sería con np.argsort

 idx = a.argsort(1)[:,:-M-1:-1] 

Aquí hay una post contiene una prueba de tiempo de ejecución que compara np.argsort y np.argpartition en un problema similar.

En caso de que alguien esté interesado tanto en los valores como en los índices correspondientes sin moderar el orden, el siguiente enfoque simple será útil. Aunque podría ser computacionalmente costoso si se trabaja con datos grandes ya que estamos usando una lista para almacenar tuplas de valor, índice.

 import numpy as np values = np.array([0.01,0.6, 0.4, 0.0, 0.1,0.7, 0.12]) # a simple array values_indices = [] # define an empty list to store values and indices while values.shape[0]>1: values_indices.append((values.max(), values.argmax())) # remove the maximum value from the array: values = np.delete(values, values.argmax()) 

El resultado final como lista de tuplas:

 values_indices [(0.7, 5), (0.6, 1), (0.4, 1), (0.12, 3), (0.1, 2), (0.01, 0)]