Como acessar a i-ésima coluna de uma matriz multidimensional NumPy?

Suponha que eu tenha:

test = numpy.array([[1, 2], [3, 4], [5, 6]])

test[i]me deixa om linha da matriz (por exemplo [1, 2]). Como posso acessar a i - ésima coluna? (por exemplo [1, 3, 5]). Além disso, isso seria uma operação cara?

>>> test[:,0]
array([1, 3, 5])

Similarmente,

>>> test[1,:]
array([3, 4])

permite acessar linhas. Isso é coberto na Seção 1.4 (Indexação) da referência NumPy . Isso é rápido, pelo menos na minha experiência. É certamente muito mais rápido do que acessar cada elemento em um loop.

E se você quiser acessar mais de uma coluna por vez, poderá:

>>> test = np.arange(9).reshape((3,3))
>>> test
array([[0, 1, 2],
       [3, 4, 5],
       [6, 7, 8]])
>>> test[:,[0,2]]
array([[0, 2],
       [3, 5],
       [6, 8]])

>>> test[:,0]
array([1, 3, 5])

esse comando fornece um vetor de linha; se você quiser fazer um loop sobre ele, tudo bem, mas se você quiser empilhar com alguma outra matriz com a dimensão 3xN, terá

ValueError: all the input arrays must have same number of dimensions

enquanto

>>> test[:,[0]]
array([[1],
       [3],
       [5]])

fornece um vetor de coluna, para que você possa executar operações de concatenação ou hstack.

por exemplo

>>> np.hstack((test, test[:,[0]]))
array([[1, 2, 1],
       [3, 4, 3],
       [5, 6, 5]])

Você também pode transpor e retornar uma linha:

In [4]: test.T[0]
Out[4]: array([1, 3, 5])

Para obter várias colunas independentes, basta:

> test[:,[0,2]]

você receberá as colunas 0 e 2

Embora a pergunta tenha sido respondida, deixe-me mencionar algumas nuances.

Digamos que você esteja interessado na primeira coluna da matriz

arr = numpy.array([[1, 2],
                   [3, 4],
                   [5, 6]])

Como você já sabe de outras respostas, para obtê-lo na forma de "vetor de linha" (matriz de formas (3,)), use o fatiamento:

arr_c1_ref = arr[:, 1]  # creates a reference to the 1st column of the arr
arr_c1_copy = arr[:, 1].copy()  # creates a copy of the 1st column of the arr

Para verificar se uma matriz é uma visualização ou cópia de outra matriz, você pode fazer o seguinte:

arr_c1_ref.base is arr  # True
arr_c1_copy.base is arr  # False

veja ndarray.base .

Além da diferença óbvia entre os dois (a modificação arr_c1_refafetará arr), o número de etapas de bytes para percorrer cada um deles é diferente:

arr_c1_ref.strides[0]  # 8 bytes
arr_c1_copy.strides[0]  # 4 bytes

ver avanços . Por que isso é importante? Imagine que você tem uma matriz muito grande em Avez de arr:

A = np.random.randint(2, size=(10000,10000), dtype='int32')
A_c1_ref = A[:, 1] 
A_c1_copy = A[:, 1].copy()

e você deseja calcular a soma de todos os elementos da primeira coluna, A_c1_ref.sum()ou seja, ou A_c1_copy.sum(). Usar a versão copiada é muito mais rápido:

%timeit A_c1_ref.sum()  # ~248 µs
%timeit A_c1_copy.sum()  # ~12.8 µs

Isso se deve ao número diferente de avanços mencionados anteriormente:

A_c1_ref.strides[0]  # 40000 bytes
A_c1_copy.strides[0]  # 4 bytes

Embora pareça que usar cópias de colunas seja melhor, nem sempre é verdade pelo fato de que fazer uma cópia leva tempo e consome mais memória (nesse caso, foram necessários aproximadamente 200 µs para criar a A_c1_copy). No entanto, se precisamos da cópia em primeiro lugar, ou precisamos executar muitas operações diferentes em uma coluna específica da matriz e estamos bem em sacrificar a memória por velocidade, fazer uma cópia é o caminho a seguir.

No caso de estarmos interessados ​​em trabalhar principalmente com colunas, pode ser uma boa ideia criar nossa matriz na ordem das colunas principais ('F') em vez da ordem das linhas principais ('C') (que é o padrão ) e faça o fatiamento como antes para obter uma coluna sem copiá-la:

A = np.asfortranarray(A)  # or np.array(A, order='F')
A_c1_ref = A[:, 1]
A_c1_ref.strides[0]  # 4 bytes
%timeit A_c1_ref.sum()  # ~12.6 µs vs ~248 µs

Agora, executar a operação de soma (ou qualquer outra) em uma exibição de coluna é muito mais rápido.

Por fim, deixe-me observar que transpor uma matriz e usar o corte de linhas é o mesmo que usar o corte de colunas na matriz original, porque a transposição é feita apenas trocando a forma e os passos da matriz original.

A.T[1,:].strides[0]  # 40000

>>> test
array([[0, 1, 2, 3, 4],
       [5, 6, 7, 8, 9]])

>>> ncol = test.shape[1]
>>> ncol
5L

Então você pode selecionar a 2ª - 4ª coluna desta maneira:

>>> test[0:, 1:(ncol - 1)]
array([[1, 2, 3],
       [6, 7, 8]])