Cómo Obtener una Columna de una Matriz en Python

Existen muchas situaciones en las que necesitamos acceder a los datos de una columna de una matriz o, incluso, separarla del resto. En este artículo te explico varias maneras de hacerlo en función del tipo de matriz que tengas que manejar y te propongo algunos ejemplos para que te quede perfectamente claro.

Una manera sencilla y compacta de extraer la columna i cuando la matriz es una lista de listas es la siguiente: columna = [fila[i] for fila in matriz]. En un array bidimensional de numpy, la forma más directa es: matriz[:,i]. En un DataFrame de pandas se obtiene así: matriz[nombre_columna].

Pero vayamos por partes viendo cada caso concreto y con ejemplos de código.

Cómo obtener columna de una lista de listas

La manera principal de crear matrices en Python es utilizando listas anidadas, es decir, crear listas de listas, por ejemplo, una matriz de 3x3, es decir, de tres filas y tres columnas la implementaríamos de esta manera:

matriz = [
	[1, 2, 3],
	[4, 5, 6],
	[7, 8, 9]
]

La solución más compacta para obtener una de las columnas de una matriz de este estilo pasa por utilizar list comprehension o comprensión de listas, que nos permite, de forma muy sencilla, generar una lista a partir de un bucle:

i = 1 #columna que queremos obtener
columna = [fila[i] for fila in matriz]

Si mostramos por pantalla la variable columna obtendremos lo siguiente, que es precisamente lo que buscamos:

[2, 5, 8]

Si prefieres no usar la comprensión de listas porque no la entiendes muy bien o porque prefieres un código más clásico (y más portable a otros lenguajes) la solución pasa por hacer un bucle más estándar, ir tomando cada valor de la columna que nos interesa y guardarlo en una lista que inicialmente está vacía.

Con la función append podemos ir añadiendo elementos al final de esa lista hasta que tengamos todos los valores de la columna:

i = 1 #columna que queremos obtener
columna = []
for fila in matriz:
	columna.append(fila[i])

Y de manera similar, con un bucle while y una variable que sirva como índice de la fila se puede hacer lo siguiente:

i = 1 #columna que queremos obtener
indice_fila = 0
columna = []
while indice_fila < len(matriz):
	columna.append(matriz[indice_fila][i])
	indice_fila += 1

Cómo obtener una columna de una lista de listas de diferentes tamaños

Hay que tener en cuenta un caso que puede suceder a veces y es que, al tener una lista de listas, es posible que cada una de las listas que representan las filas tenga un tamaño diferente de forma que lo que tenemos es una matriz irregular. Por ejemplo:

matriz_irregular = [
	[1, 2, 3],
	[4, 5],
	[6, 7, 8, 9]
]

Es bastante evidente que si aplicamos la solución anterior a este ejemplo obtendremos un error de índice fuera de rango si intentamos obtener, por ejemplo, la columna 2, ya que la segunda fila solo tiene dos elementos (columnas 0 y 1).

Para subsanar esta situación tendremos que decidir qué es lo que nos interesa más, pues tenemos dos opciones:

1. Incluir un valor None en la lista para aquellas filas que no tienen elementos suficientes.
2. Ignorar las filas que no tienen elementos suficientes, de forma que la lista que obtengamos tenga menos valores que el número de filas de la matriz.

Si optamos por la primera opción podemos alterar la solución compacta, mostrada más arriba, para incluir un condicional que comprueba si la fila tiene valores suficientes y si no los tiene inserta el valor None:

i = 2 #columna que queremos obtener
columna = [fila[i] if i < len(fila) else None for fila in matriz_irregular]

Si preferimos o necesitamos la segunda opción, de manera similar podemos hacer lo siguiente:

i = 2
columna = [fila[i] for fila in matriz if i < len(fila)]

Cómo obtener columna de una matriz de NumPy

Si necesitamos hacer operaciones matemáticas con matrices o trabajar de manera más científica, es recomendable utilizar la librería NumPy. Asegúrate, para ello, de tener instalado el módulo numpy mediante el comando pip install numpy.

La forma de manejar matrices con este módulo se basa en el uso de arrays que es un tipo de dato propio de esta librería. De esta manera, podemos definir una matriz como sigue:

import numpy as np

matriz = [
	[21, 22, 23],
	[34, 35, 36],
	[47, 48, 49]
]
matriz_np = np.array(matriz)

Una vez con la matriz definida es sumamente fácil obtener una de sus columnas. Se puede hacer de la siguiente manera:

i = 1 #columna que queremos obtener
columna = matriz_np[:,i]

Fíjate en que ahora utilizamos dos índices para acceder a los elementos de la matriz, entre corchetes y separados por una coma. Utilizamos la notación de dos puntos (:) para acceder a todas las filas como primer índice, y como segundo índice indicamos la columna a obtener. Si ahora mostramos por pantalla el contenido de la variable columna obtendremos lo siguiente, que es precisamente el contenido de la columna 1:

[22 35 48]

Aunque podemos definirlas, NumPy no permite trabajar de manera tan sencilla con matrices de filas de diferentes longitudes por lo que tendríamos que optar por una solución similar a la presentada previamente en el caso de tener una lista de listas. En lugar de hacer una solución compacta como la de antes, que también se podría plantear, en esa ocasión hacemos un bucle más estándar (para que queden bien claras ambas alternativas):

import numpy as np

matriz_irregular = [
	[21, 22, 23, 24],
	[34, 35],
	[47, 48, 49]
]
matriz_irregular_np = np.array(matriz_irregular)

i = 2 #columna que queremos obtener
columna = []
for fila in matriz_irregular_np:
	if i < len(fila):
		columna.append(fila[i])

Cómo obtener una columna de un DataFrame de pandas

pandas es la librería de referencia para trabajar con datos en Python. El principal tipo de dato propio de pandas es el DataFrame, que representa una matriz o tabla de datos.

Para no extenderme demasiado aquí, te dejo un artículo en el que te cuento cómo crear un DataFrame a partir de diversas fuentes de datos con pandas.

Dado que la dimensión principal de un DataFrame son las columnas y no las filas, acceder a los valores de una columna es extremadamente sencillo.

En el siguiente ejemplo te muestro cómo acceder a diferentes columnas de la matriz representada mediante un DataFrame. Basta con utilizar la notación de corchetes con un único índice que es el nombre de la columna que queremos obtener:

import pandas as pd

# preparamos los datos como una lista de listas
datos = [
    [10, 20, 30, 40],
    [11, 21, 31, 41],
    [12, 22, 32, 42]
]

# creamos el DataFrame indicando los nombres de las columnas
matriz = pd.DataFrame(datos, columns=['A', 'B', 'C', 'D'])

print(matriz) # mostramos la matriz por pantalla

columna_a = matriz['A'] # accedemos a la columna A
columna_c = matriz['C'] # accedemos a la columa C

print(columna_a.values) # mostramos los valores de la columna A

datos_c = list(columna_c) # convertimos la columna C en una lista

print(datos_c) # mostramos la columna C

Si te fijas en las líneas 15 y 16 el acceso a las columnas es muy claro y directo, indicando entre corchetes el nombre de la columna.

Ten en cuenta que esto nos permite obtener la columna en un formato de datos propio de pandas, que es Series, que en este caso representa una lista de elementos. No obstante, podemos acceder directamente a sus valores a través de la propiedad values (línea 18) o convertirla directamente a una lista (línea 20) con la función list.

El resultado por pantalla será el siguiente:

    A   B   C   D
 0  10  20  30  40
 1  11  21  31  41
 2  12  22  32  42
 [10 11 12]
 [30, 31, 32]

Cómo eliminar una columna de una matriz

Si en lugar de obtener los valores de una columna lo que pretendemos es eliminar la columna de una matriz la cosa cambia un poco. Puedes apoyarte para esto en las funciones pop y del que nos permiten borrar un elemento en una posición concreta de una lista.

Si tenemos una lista de listas que forma una matriz regular donde todas las filas tienen el mismo tamaño, basta con recorrer cada fila de la matriz borrando de cada una el elemento deseado, que es el que corresponde con la columna a eliminar.

Alternativa con la función del:

i = 1 #columna que queremos eliminar
for fila in matriz:
	del fila[i]

Y, de la misma manera, con la función pop quedaría así:

i = 1 #columna que queremos eliminar
for fila in matriz:
	fila.pop(i)

Si la matriz que tenemos es una lista de listas donde cada fila tiene un tamaño diferente, tendremos que alterar el bucle para incluir una condición comprobando que el elemento a borrar exista en cada fila. De no hacer esto, tanto el uso de pop como de del generarán en error:

i = 1 #columna que queremos eliminar
for fila in matriz:
	if i < len(fila):
		del fila[i]

Si estamos trabajando con una matriz de numpy, para eliminar una columna en concreto podemos usar la función delete a la que hay que indicarle tres parámetros. El primero es la matriz original, el segundo el índice de la columna (o fila) que queremos eliminar y, finalmente, el parámetro axis, o eje, con el que indicamos la dimensión en la que haremos el borrado (si fuera 0 borraríamos una fila, si es 1, una columna).

import numpy as np

matriz = [
	[21, 22, 23],
	[34, 35, 36],
	[47, 48, 49]
]
matriz_np = np.array(matriz)

i = 1 #columna que queremos eliminar
matriz_nueva_np = np.delete(matriz_np, i, axis=1)

Si mostramos la variable matriz_nueva_np por pantalla obtenemos lo siguiente:

[[21 23]
 [34 36]
 [47 49]]

En el caso de querer borrar una columna de un DataFrame de pandas tenemos tres opciones:

1. Usar la palabra clave del cómo hemos visto antes.
2. Usar la función pop del DataFrame, que elimina la columna.
3. Usar la función drop del DataFrame, que devuelve un nuevo DataFrame sin la columna indicada. Al igual que sucede con el delete de NumPy, tendremos que indicarle el eje en el que queremos eliminar mediante el parámetro axis, proporcionándole un 1.

Veamos las tres opciones en un único ejemplo:

import pandas as pd

# preparamos los datos como una lista de listas
datos = [
    [10, 20, 30, 40],
    [11, 21, 31, 41],
    [12, 22, 32, 42]
]

# creamos el DataFrame indicando los nombres de las columnas
matriz = pd.DataFrame(datos, columns=['A', 'B', 'C', 'D'])

del matriz['A']
matriz.pop('C')
matriz = matriz.drop('D', axis=1)

print(matriz)

En las líneas 13, 14 y 15 eliminamos tres columnas diferentes con cada uno de los tres métodos. Como puedes ver, al DataFrame solo le queda una columna cuando se muestra por pantalla:

    B
 0  20
 1  21
 2  22

Ejemplos de acceso a columnas en matrices

Para contextualizar un poco y para terminar de aclarar estos conceptos te propongo un ejemplo que podría ser real.

Vamos a imaginar que tenemos una matriz de 10 filas y de 7 columnas. Cada fila representa una semana y cada columna representa un día de la semana de lunes a domingo, donde la columna 0 representa el lunes, el 1 el martes, etc.

Cada valor de la matriz es la temperatura más alta registrada para una semana en un día concreto en una estación meteorológica de Madrid. Es decir, que en la posición [0][0] tenemos la temperatura más alta registrada el lunes en la primera de las 10 semanas y en la [4][5] tenemos la temperatura más alta registrada el sábado de la quinta semana.

Nos piden que calculemos la media de las temperaturas más altas recogidas el fin de semana durante esas 10 semanas.

Una manera sencilla de realizar este cálculo es obteniendo las columnas de la matriz que corresponden al sábado y al domingo y calculando el valor medio de todos los valores contenidos en esas columnas. A continuación te muestro el código de todo el ejemplo:

#Esta es la matriz de temperaturas, lo normal sería leerlas de un fichero
temperaturas = [
	[12.5, 14.8, 13.5, 13.5, 12.8, 13.4, 13.1],
	[13.3, 15.7, 11.5, 12.7, 13.7, 14.4, 14.4],
	[14.7, 15.6, 10.5, 11.5, 14.6, 14.0, 13.3],
	[13.4, 14.2, 13.5, 10.5, 16.8, 14.6, 12.1],
	[14.9, 15.9, 12.5, 11.6, 18.7, 16.7, 11.6],
	[15.4, 15.0, 13.5, 12.6, 20.1, 13.8, 12.7],
	[16.1, 16.0, 13.5, 11.2, 17.8, 12.1, 13.9],
	[14.1, 13.0, 14.5, 11.1, 15.1, 12.0, 13.2],
	[13.2, 12.4, 11.5, 12.0, 14.0, 11.8, 14.0],
	[13.7, 11.3, 11.5, 15.0, 13.1, 13.3, 15.1],
]

#Preparamos los datos necesarios
indice_sabado = 5
indice_domingo = 6

#Obtenemos las columnas
columna_sabado = [semana[indice_sabado] for semana in temperaturas]
columna_domingo = [semana[indice_domingo] for semana in temperaturas]

#Mostramos los valores de las columnas por pantalla
print(f'Temperaturas de los sábados: {columna_sabado}')
print(f'Temperaturas de los domingos: {columna_domingo}')

#Creamos una única lista con todos los valores del fin de semana
temperaturas_fin_de_semana = columna_sabado + columna_domingo

#Calculamos el valor medio sumando todos los valores
#y dividiendo entre el número de valores
temperatura_media = sum(temperaturas_fin_de_semana) / len(temperaturas_fin_de_semana)

#Finalmente mostramos la temperatura media del fin de semana
print(f'La media de las temperaturas más altas para el fin de semana es de {temperatura_media} grados cetígrados.')

Al ejecutar este programa obtenemos como resultado:

Temperaturas de los sábados: [13.4, 14.4, 14.0, 14.6, 16.7, 13.8, 12.1, 12.0, 11.8, 13.3]
Temperaturas de los domingos: [13.1, 14.4, 13.3, 12.1, 11.6, 12.7, 13.9, 13.2, 14.0, 15.1]
La media de las temperaturas más altas para el fin de semana es de 13.475 grados centígrados.

Consejos y buenas prácticas

Si en tu trabajo o programa tienes que obtener o eliminar columnas frecuentemente de matrices, te recomiendo que desarrolles tus propias funciones con lo que aquí te he explicado para facilitarte la tarea. Además, así evitas repetir código una y otra vez.

Vamos a suponer que necesitamos tanto obtener columnas como eliminarlas y tenemos que trabajar con matrices hechas con una lista de listas y además las filas pueden tener longitudes diferentes. Para esto desarrollamos dos funciones, una para obtener una columna y otra para eliminar una columna.

def obtener_columna(matriz, indice_columna):
	return [fila[indice_columna] if indice_columna < len(fila) else None for fila in matriz]

def eliminar_columna(matriz, indice_columna):
	for fila in matriz:
		if i < len(fila):
			del fila[indice_columna]

Así, cada vez que queramos obtener una columna invocamos a la función obtener_columna indicando de qué matriz vamos a obtenerla y el índice de la columna en cuestión. Para eliminar una columna llamamos a la función eliminar_columna indicando los mismos parámetros. Por ejemplo, y suponiendo la variable matriz:

#obtenemos la columna 1
columna_1 = obtener_columna(matriz, 1)

#obtenemos la columna 7
columna_7 = obtener_columna(matriz, 7)

#borramos la columna 5
eliminar_columna(matriz, 5)

Espero que este artículo te haya sido de ayuda. Compártelo con los botones que encontrarás abajo, ¡gracias!