Como obter uma coluna de um array NumPy

Para selecionar uma coluna de um array NumPy bidimensional, use array[:, column_index]. Os dois-pontos selecionam todas as linhas, e o número inteiro seleciona uma coluna com índice iniciado em zero. Por exemplo, data[:, 1] retorna a segunda coluna como um array unidimensional. Em vez disso, use data[:, 1:2] quando o código subsequente exigir um resultado bidimensional com formato (rows, 1).

Selecionar uma coluna e controlar seu formato

O NumPy indexa um array 2D como array[row_selection, column_selection]. Omitir a vírgula muda a operação: data[:][1] primeiro cria um slice completo e depois seleciona a linha 1, portanto não seleciona a coluna 1.

O exemplo a seguir abrange os casos comuns de seleção de colunas em um único programa reproduzível:

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)

Resultado:

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]

Um índice de coluna inteiro remove esse eixo, o que explica o formato (3,) de second. Um slice de um único elemento mantém o eixo, produzindo o formato (3, 1) de second_2d. Essa distinção é importante quando uma API espera uma entrada semelhante a uma matriz em vez de um vetor simples.

Selecionar várias colunas

Use um slice como data[:, :2] para colunas adjacentes. O slicing expressa um intervalo contínuo e geralmente produz uma visualização. Para colunas não adjacentes ou uma ordem personalizada, forneça uma lista de índices: data[:, [2, 0]]. Isso é indexação avançada e produz uma cópia em vez de uma visualização.

Uma lista de índices sempre mantém o eixo das colunas. Assim, data[:, [1]] tem o formato (3, 1), enquanto data[:, 1] tem o formato (3,). Uma lista também pode repetir uma coluna ou alterar sua posição, portanto valide índices fornecidos dinamicamente quando duplicatas ou a ordenação forem inesperadas. Use um slice quando as colunas formarem um intervalo simples e uma lista de índices quando a própria seleção for uma sequência.

np.take é uma alternativa que explicita o eixo. np.take(data, [0, 2], axis=1) seleciona as colunas 0 e 2; axis=1 é essencial porque as colunas ocupam o segundo eixo. Ela é particularmente útil quando o eixo ou os índices estão armazenados em variáveis.

Filtrar linhas antes de ler uma coluna

Uma condição booleana normalmente identifica linhas, não colunas. Crie a máscara de linhas a partir de uma coluna e coloque a coluna de saída desejada depois da vírgula. No exemplo, data[:, 0] >= 20 mantém as duas últimas linhas, e data[data[:, 0] >= 20, 2] retorna [22 32] da coluna 2. O comprimento da máscara deve corresponder ao número de linhas.

Decidir se você precisa de uma visualização ou de uma cópia

De acordo com as regras de indexação do NumPy, o slicing básico retorna uma visualização quando possível, enquanto a indexação avançada retorna uma cópia. Consequentemente, atribuir por meio de view = data[:, 1] pode alterar data, como demonstra o 99 na saída. Chame .copy() quando a coluna extraída precisar ser independente. Os valores copiados permanecem [11 21 31] depois que a visualização altera o array de origem.

Essas visualizações também mantêm acessível a memória de seu array base. Copie uma coluna pequena se manter um array de origem grande desperdiçaria memória depois que o restante dessa origem não fosse mais necessário.

Verificar as dimensões e os limites das colunas

Essas expressões pressupõem um array 2D. Um array 1D não tem eixo de colunas, portanto array[:, 0] gera um erro de indexação. Verifique array.ndim == 2 quando o formato da entrada for incerto. Para um array 2D, os índices de coluna não negativos válidos vão de 0 a array.shape[1] - 1; um índice fora desse intervalo gera IndexError. Índices negativos também são válidos dentro dos limites, portanto array[:, -1] seleciona a última coluna.