Cómo obtener una columna de un array de NumPy

Para seleccionar una columna de un array bidimensional de NumPy, utiliza array[:, column_index]. Los dos puntos seleccionan todas las filas y el entero selecciona una columna mediante un índice que empieza en cero. Por ejemplo, data[:, 1] devuelve la segunda columna como un array unidimensional. En cambio, utiliza data[:, 1:2] cuando el código posterior requiera un resultado bidimensional con forma (rows, 1).

Seleccionar una columna y controlar su forma

NumPy indexa un array bidimensional como array[row_selection, column_selection]. Omitir la coma cambia la operación: data[:][1] crea primero un slice completo y después selecciona la fila 1, por lo que no selecciona la columna 1.

El ejemplo siguiente reúne los casos habituales de selección de columnas en un programa reproducible:

import numpy as np

data = np.array([[10, 11, 12], [20, 21, 22], [30, 31, 32]])

second = data[:, 1]
second_2d = data[:, 1:2]
adjacent = data[:, :2]
reordered = data[:, [2, 0]]
taken = np.take(data, [0, 2], axis=1)
filtered = data[data[:, 0] >= 20, 2]

print("second:", second, second.shape)
print("second_2d:\n", second_2d, second_2d.shape)
print("adjacent:\n", adjacent)
print("reordered:\n", reordered)
print("taken:\n", taken)
print("filtered:", filtered)

view = data[:, 1]
copy = data[:, 1].copy()
view[0] = 99
print("data after view edit:\n", data)
print("copy:", copy)

Salida:

second: [11 21 31] (3,)
second_2d:
 [[11]
 [21]
 [31]] (3, 1)
adjacent:
 [[10 11]
 [20 21]
 [30 31]]
reordered:
 [[12 10]
 [22 20]
 [32 30]]
taken:
 [[10 12]
 [20 22]
 [30 32]]
filtered: [22 32]
data after view edit:
 [[10 99 12]
 [20 21 22]
 [30 31 32]]
copy: [11 21 31]

Un índice de columna entero elimina ese eje, lo que explica la forma (3,) de second. Un slice de un solo elemento conserva el eje y produce la forma (3, 1) de second_2d. Esta diferencia es importante cuando una API espera una entrada similar a una matriz en lugar de un vector plano.

Seleccionar varias columnas

Utiliza un slice como data[:, :2] para seleccionar columnas adyacentes. El slicing expresa un intervalo continuo y normalmente produce una vista. Para seleccionar columnas no adyacentes o establecer un orden personalizado, proporciona una lista de índices: data[:, [2, 0]]. Esto es indexación avanzada y produce una copia en lugar de una vista.

Una lista de índices siempre conserva el eje de columnas. Por tanto, data[:, [1]] tiene la forma (3, 1), mientras que data[:, 1] tiene la forma (3,). Una lista también puede repetir una columna o cambiar su posición, así que valida los índices proporcionados dinámicamente si los duplicados o el orden pudieran resultar inesperados. Utiliza un slice cuando las columnas formen un intervalo sencillo y una lista de índices cuando la selección sea una secuencia.

np.take es una alternativa que hace explícito el eje. np.take(data, [0, 2], axis=1) selecciona las columnas 0 y 2; axis=1 es imprescindible porque las columnas ocupan el segundo eje. Resulta especialmente útil cuando el eje o los índices están almacenados en variables.

Filtrar filas antes de leer una columna

Una condición booleana normalmente identifica filas, no columnas. Crea la máscara de filas a partir de una columna y coloca después de la coma la columna que quieres obtener. En el ejemplo, data[:, 0] >= 20 conserva las dos últimas filas, y data[data[:, 0] >= 20, 2] devuelve [22 32] de la columna 2. La longitud de la máscara debe coincidir con el número de filas.

Decidir si necesitas una vista o una copia

Según las reglas de indexación de NumPy, el slicing básico devuelve una vista siempre que sea posible, mientras que la indexación avanzada devuelve una copia. Por ello, asignar mediante view = data[:, 1] puede modificar data, como demuestra el 99 de la salida. Llama a .copy() cuando la columna extraída deba ser independiente. Los valores copiados permanecen como [11 21 31] después de que la vista modifique el array de origen.

Estas vistas también mantienen accesible la memoria de su array base. Copia una columna pequeña si conservar un array de origen grande desperdiciaría memoria una vez que ya no se necesite el resto de ese array.

Comprobar las dimensiones y los límites de las columnas

Estas expresiones presuponen un array bidimensional. Un array unidimensional no tiene eje de columnas, por lo que array[:, 0] genera un error de indexación. Comprueba array.ndim == 2 cuando no conozcas con certeza la forma de la entrada. En un array bidimensional, los índices de columna no negativos válidos van de 0 a array.shape[1] - 1; un índice fuera de ese intervalo genera IndexError. Los índices negativos también son válidos dentro de los límites, por lo que array[:, -1] selecciona la última columna.