Obtener los índices de los N valores máximos en NumPy

En NumPy existe la función np.max() para obtener el máximo de un vector o matriz y np.argmax() para obtener la posición del máximo. De forma análoga también existen las funciones np.min() y np.argmin() para el caso de querer obtener el mínimo o su posición. Pero, en el caso de que se desee obtener los índices de los N valores máximos en NumPy no existe una función directa y es necesario hacer unos pasos más.

Posición del máximo de un elemento en NumPy

Como se ha comentado en la introducción, para obtener el máximo de un elemento en vector o matriz de NumPy se puede recurrir a la función np.max() y en el caso de querer su índice se puede recurrir a np.argmax(). Lo que se muestra en el siguiente ejemplo.

import numpy as np

# Fijar la semilla para los números aleatorios
np.random.seed(1)

# Crear vector NumPy aleatorio
arr = np.random.randint(1, 20, 12)

# Encontrar el valor máximo y su posición
max_val = np.max(arr)
max_idx = np.argmax(arr)

print("Vector: ", arr)
print("Valor máximo: ", max_val)
print("Posición del valor máximo: ", max_idx)

Vector:  [ 6 12 13  9 10 12  6 16  1 17  2 13]
Valor máximo:  17
Posición del valor máximo:  9

En este código lo primero que se hace es importar NumPy y fijar la semilla a 1 para garantizar que siempre se obtengan los mismos números aleatorios siempre que se ejecute el código. Una vez hecho esto se generan 12 números aleatorios enteros entre 1 y 20 con la función np.random.randint(). Tras la obtención del vector obtener el máximo, la posición de este y se muestran por pantalla. Como se puede ver en los resultados el valor máximo es 17 y este se encuentra en el índice 9.

Índices de los N valores máximos en NumPy

Una vez que se ha obtenido el valor máximo, para conseguir la posición de los N valores máximos se podría generar una copia ordenada de elementos del vector y buscar su posición en el original. Pero este no es un método eficiente. La mejor opción es recurrir a la función argsort() que devuelve los índices que ordena el vector. Así, para conseguir los índices de los N valores máximos solamente se tendría que obtener los 5 últimos. Proceso que se muestra en el siguiente ejemplo de código.

# Obtener los índices de los 5 valores máximos
max_5_idx = np.argsort(arr)[-5:]

print("Vector: ", arr)
print("Valor los máximos: ", arr[max_5_idx])
print("Posición del los máximos: ", max_5_idx)

Vector:  [ 6 12 13  9 10 12  6 16  1 17  2 13]
Valor los máximos:  [12 13 13 16 17]
Posición del los máximos:  [ 5  2 11  7  9]

En este caso lo único que se hace es llamar a la función np.argsort(arr) y asignar a la variable max_5_idx la posición de los 5 valores máximos. Una vez hecho esto solamente se muestran los resultados por pantalla.

Índices de los N valores mínimos en NumPy

En el caso de querer obtener los mínimos en lugar de los máximos solamente se tienen que cambiar las funciones con max por la versión con min y, en el caso de argsort(), obtener los primeros elementos en lugar de los últimos. Así se puede repetir los ejercicios anteriores para los mínimos usando el siguiente código.

# Encontrar los valores mínimos y su posiciones
min_val = np.min(arr)
min_idx = np.argmin(arr)
min_5_idx = np.argsort(arr,)[:5]

print("Vector: ", arr)
print("Valor mínimo: ", min_val)
print("Posición del valor mínimo: ", min_idx)
print("Valor los mínimos: ", arr[min_5_idx])
print("Posición del los mínimos: ", min_5_idx)

Vector:  [ 6 12 13  9 10 12  6 16  1 17  2 13]
Valor mínimo:  1
Posición del valor mínimo:  8
Valor los mínimos:  [1 2 6 6 9]
Posición del los mínimos:  [ 8 10  0  6  3]

Conclusiones

La función np.argsort() devuelve los índices con los que se puede ordenar un matriz, por lo que se puede usar para conseguir los índices de los N valores máximos en NumPy, también para el mínimo. Así, en el caso de que sea necesario obtener los N valores extremos de una muestra no es necesario recurrir a bucles que son menos eficientes.