NumPy: Ordenar en base a la suma de las columnas o filas de matrices NumPy

El año pasado publiqué una entrada en la que se enseñaba a ordenar las matrices de NumPy en base a los valores de una fila o columna. Para lo que se usaba el método argsort(). En esta ocasión veremos otro problema similar y bastante habitual: ordenar en base a la suma de las columnas o de las filas.

Obtener la suma de las columnas o filas en NumPy

Para continuar es necesario crear primero una matriz, para lo que se puede usar simplemente el siguiente código

import numpy as np

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

Ahora, la suma de los elementos de esta matriz se puede obtener mediante el método sum() que tienen los objetos ndarray. Un método que si no se indica nada devuelve la suma de todos los elementos, pero, mediante la propiedad axis, es posible indicarle que se desea la suma de las filas (1) o columnas (0). Por ejemplo, para obtener la suma de las columnas se puede escribir

arr.sum(axis=0)

array([64, 39, 67])

Ordenar en base a la suma de columnas

Ahora bien, para reordenar la matriz en base a estos valores es necesario obtener los índices que ordenaría este vector, lo que nos proporciona el método argsort(). Así se pueden obtener los índices mediante la línea

arr.sum(axis=0).argsort()

array([1, 0, 2])

Donde se puede ver que en la segunda posición se encuentra la columna cuya suma es la menor de las tres, en la primera el valor intermedio y en la última con la mayor suma. Esto es, los índices con los que se podría ordenar la matriz en orden ascendente.

arr.sum(axis=0)[arr.sum(axis=0).argsort()]

array([39, 64, 67])

Lo que también se puede emplear para ordenar las columnas de la matriz en base a la suma de estas. Algo que se puede, tal como se explicó anteriormente, se puede obtener simplemente con el siguiente código.

arr[:, arr.sum(axis=0).argsort()]

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

Ordenar en orden decreciente

Desafortunadamente, el método argsort() no dispone de una propiedad con la que se obtengan los índices que ordenen el vector original en orden descendente. Aunque esto es algo que se puede solucionar fácilmente mediante el uso del operador :, simplemente invirtiendo el orden del array tal como se muestra a continuación.

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

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

Ordenar en base a la suma de las filas

Finalmente, para ordenar en base a la suma de las filas solamente hay que cambiar el eje sobre el que suma el método sum(), obtener los indices correspondientes y aplicarlos para ordenar las filas. Así para ordenar en base a la suma de las filas solamente hay que escribir.

arr[arr.sum(axis=0).argsort()]

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

Conclusiones

En esta ocasión hemos visto cómo se puede ordenar en base a la suma de las columnas o filas la matriz NumPy. Comprobado como una operación que es bastante habitual se puede realizar fácilmente en Python combinando el uso de los métodos sum() y argsort().