NumPy: Obtener los índices de los máximos de un vector

Para ordenar los elementos de un vector de NumPy existe el método sort(). Pero, en algunas ocasiones, puede que esto no sea los necesario, sino los índices de los máximos o mínimos de un vector. Por ejemplo, para seleccionar los elementos asociados en otro vector. En estas ocasiones se puede recurrir al método argsort() con el cual se puede obtener justamente este resultado.

El método argsort()

En NumPy el método argsort() tiene la siguiente forma

np.argsort(arr, axis=-1, kind=None, order=None)

en donde:

• arr: un array o matriz compatible con NumPy
• axis: el eje sobre el que se realiza la ordenación, por defecto es el ultimo eje, -1, y si se indica None la matriz se ordenará como si fuese un vector.
• kind: una cadena de texto con el método de ordenación que se desea emplear.
• order: solamente se usa en vectores con campos definidos para indicar el orden por el que se ordenan los elementos.

El método también está disponible como propiedad en todas las matrices de NumPy, donde no es necesario indicar la matriz.

Uso básico de argsort()

Al llamar a argsort() se obtiene un vector con los indices que ordenaría los elementos de este en orden ascendente. Lo que se puede usar para ordenar un vector tal como se muestra a continuación

import numpy as np

arr = np.array([7, 1, 6, 3, 9, 2, 4, 5, 8])

arr.argsort()      # [1, 5, 3, 6, 7, 2, 0, 8, 4]
arr[arr.argsort()] # [1, 2, 3, 4, 5, 6, 7, 8, 9]

En este ejemplo se puede ver que el método argsort() indica que el primer elemento del vector ordenado es aquel que se sitúa en la posición correspondiente al índice 1, que en este caso es 1. El segundo es el situado en índice 5, el cual se puede comprobar que es el 5. Así hasta el último que el 4 correspondiente al 9.

Otra cosa que se puede ver en el ejemplo es al usar el resultado en el vector original se consigue esté ordenado en orden ascendente.

Obtener los índices de los máximos del vector

El método únicamente devuelve los índices que ordenan el vector en orden ascendente, no existe actualmente una opción con la que se pueda cambiar el orden. Aunque esto es algo que se puede solucionar de una manera sencilla: aplicando el método a un vector con el signo cambiado. Con lo que es posible ordenar el vector en orden ascendente tal como se muestra en el siguiente ejemplo.

(-arr).argsort()      # [4, 8, 0, 2, 7, 6, 3, 5, 1]
arr[(-arr).argsort()] # [9, 8, 7, 6, 5, 4, 3, 2, 1]

Ahora, en caso de que solamente sean necesarios los primeros N índices no hay más que seleccionar estos. Lo que se puede hacer tanto para obtener los máximos como los mínimos. Algo que se puede ver en el siguiente ejemplo.

Conclusiones

Hoy hemos visto cómo se puede usar el método argsort() de NumPy para obtener los índices con los que ordenar los elementos de un vector. A pesar de que solamente devuelve los índices para producir el resultado en orden ascendente, con un pequeño truco se pueden obtener estos en orden descendente. Un método que también se puede utilizar para ordenar matrices por filas o columnas o en base a operaciones como sumas.