NumPy: Ordenar matrices de NumPy por fila o columna

Cuando estamos trabajando en NumPy con matrices, arrays 2D, nos podemos plantear cómo reordenar estas en base a los valores de las filas o las columnas. Una operación que se puede realizar fácilmente utilizando el método argsort() presente en los array de NumPy. En esta entrada veremos los métodos para ordenar matrices de NumPy por fila o columna.

Creación de una matriz

En esta ocasión necesitamos una matriz sencilla que se puede crear manualmente. Para evitar confusión lo mejor es crear una que no tenga ninguna fila o columna ordenada, como la que se puede crear con el siguiente comando.

import numpy as np

arr = np.array([[21, 12, 33],
                [12,  7, 13],
                [31, 20, 21]])

Ordenar en base a los valores de una columna

Ahora supongamos que queremos ordenar los valores de las filas en base a los de la primera columna. Para lo que solamente deberíamos seleccionar esta y utilizar el método argsort(). Un método que devuelve los índices que ordenarían el vector en orden ascendente. Justamente lo que necesitamos para poder ordenar la matriz en base a los valores de la columna seleccionada. Así, empleando el resultado del método se puede reordenar la matriz original en base a los valores de una columna, en nuestro ejemplo la primera.

arr[arr[:,0].argsort()]

array([[12,  8, 13, 31],
       [21, 12, 33, 22],
       [31, 20, 21,  7]])

La matriz resultante se encuentra ordenada en orden ascendente. Para ordenarla en orden descendente solamente se le debe dar la vuelta a los índices antes de usarlos en la matriz. Algo que se puede hacer fácilmente con operador dos puntos (:) listando los elementos del final al principio [::-1]. Así es posible invertir el orden mediante el siguiente código.

arr[arr[:,0].argsort()[::-1]]

array([[31, 20, 21,  7],
       [21, 12, 33, 22],
       [12,  8, 13, 31]])

Ordenar en base a los valores de una fila

En el caso de querer ordenar por filas debemos seguir el mismo método que para las columnas. Solo que en este caso el vector a extraer debería ser una fila. Al mismo tiempo que los índices obtenidos mediante el método argsort() se aplican para ordenar las columnas. Lo que se puede conseguir mediante un código como el siguiente.

arr[:, arr[0].argsort()]

array([[12, 21, 22, 33],
       [ 8, 12, 31, 13],
       [20, 31,  7, 21]])

En este caso los datos también son ordenados de forma ascendente. Para ordenarlos en orden descendente también se puede invertir los índices antes de reordenar la matriz.

Conclusiones

En esta entrada hemos visto cómo se puede emplear el método argsort() para ordenar matrices de NumPy por fila o columna. Tarea que es bastante habitual.