Comment accéder à la ième colonne d'un tableau multidimensionnel NumPy?

Supposons que j'ai:

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

test[i]me donne la ième ligne du tableau (par exemple [1, 2]). Comment puis-je accéder à la ième colonne? (par exemple [1, 3, 5]). Est-ce aussi une opération coûteuse?

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

De même,

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

vous permet d'accéder aux lignes. Ceci est couvert dans la section 1.4 (Indexation) de la référence NumPy . C'est rapide, du moins d'après mon expérience. C'est certainement beaucoup plus rapide que d'accéder à chaque élément d'une boucle.

Et si vous souhaitez accéder à plusieurs colonnes à la fois, vous pouvez le faire:

>>> 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])

cette commande vous donne un vecteur de ligne, si vous voulez simplement faire une boucle dessus, c'est bien, mais si vous voulez empiler avec un autre tableau de dimension 3xN, vous aurez

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

tandis que

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

vous donne un vecteur de colonne, de sorte que vous pouvez faire une opération de concaténation ou de pile d'empilement.

par exemple

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

Vous pouvez également transposer et renvoyer une ligne:

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

Pour obtenir plusieurs colonnes indépendantes, il suffit de:

> test[:,[0,2]]

vous obtiendrez des colonnes 0 et 2

Bien que la question ait été répondue, permettez-moi de mentionner quelques nuances.

Disons que vous êtes intéressé par la première colonne du tableau

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

Comme vous le savez déjà dans d'autres réponses, pour l'obtenir sous forme de "vecteur ligne" (tableau de forme (3,)), vous utilisez le découpage:

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

Pour vérifier si un tableau est une vue ou une copie d'un autre tableau, vous pouvez procéder comme suit:

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

voir ndarray.base .

Outre la différence évidente entre les deux (la modification arr_c1_refaffectera arr), le nombre d'étapes d'octet pour parcourir chacune d'entre elles est différent:

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

voir les foulées . Pourquoi est-ce important? Imaginez que vous ayez un très grand tableau Aau lieu de arr:

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

et vous voulez calculer la somme de tous les éléments de la première colonne, c'est A_c1_ref.sum()-à- dire ou A_c1_copy.sum(). L'utilisation de la version copiée est beaucoup plus rapide:

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

Cela est dû au nombre différent de pas mentionnés précédemment:

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

Bien qu'il puisse sembler préférable d'utiliser des copies de colonne, ce n'est pas toujours le cas, car la création d'une copie prend du temps et utilise plus de mémoire (dans ce cas, il m'a fallu environ 200 µs pour créer le A_c1_copy ). Cependant, si nous avons besoin de la copie en premier lieu, ou si nous devons effectuer de nombreuses opérations différentes sur une colonne spécifique du tableau et que nous acceptons de sacrifier la mémoire pour la vitesse, alors faire une copie est la voie à suivre.

Dans le cas où nous souhaiterions travailler principalement avec des colonnes, il pourrait être judicieux de créer notre tableau dans l'ordre des colonnes principales ('F') au lieu de l'ordre des lignes principales ('C') (qui est la valeur par défaut). ), puis effectuez le découpage comme précédemment pour obtenir une colonne sans la copier:

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

Maintenant, effectuer l'opération de somme (ou toute autre) sur une vue en colonne est beaucoup plus rapide.

Enfin, permettez-moi de noter que la transposition d'un tableau et l'utilisation du découpage en ligne sont les mêmes que l'utilisation du découpage en colonne sur le tableau d'origine, car la transposition se fait en échangeant simplement la forme et les enjambées du tableau d'origine.

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

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

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

Ensuite, vous pouvez sélectionner la 2ème - 4ème colonne de cette façon:

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