web-dev-qa-db-fra.com

NumPy sélection d'un index de colonne spécifique par ligne à l'aide d'une liste d'index

J'ai du mal à sélectionner les colonnes spécifiques par ligne d'une matrice NumPy.

Supposons que j'ai la matrice suivante que j'appellerais X:

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

J'ai également un list d'index de colonne pour chaque ligne que j'appellerais Y:

[1, 0, 2]

J'ai besoin d'obtenir les valeurs:

[2]
[4]
[9]

Au lieu d'un list avec des index Y, je peux également produire une matrice avec la même forme que X où chaque colonne est un bool/int dans la plage 0-1, indiquant s'il s'agit de la colonne requise.

[0, 1, 0]
[1, 0, 0]
[0, 0, 1]

Je sais que cela peut être fait en itérant sur le tableau et en sélectionnant les valeurs de colonne dont j'ai besoin. Cependant, cela sera exécuté fréquemment sur de grands tableaux de données et c'est pourquoi il doit fonctionner aussi vite que possible.

je me demandais donc s'il y avait une meilleure solution?

Merci.

pythonpython-2.7numpy
68
Zee

Si vous avez un tableau booléen, vous pouvez faire une sélection directe basée sur cela comme ceci:

>>> a = np.array([True, True, True, False, False])
>>> b = np.array([1,2,3,4,5])
>>> b[a]
array([1, 2, 3])

Pour suivre votre exemple initial, vous pouvez procéder comme suit:

>>> a = np.array([[1,2,3], [4,5,6], [7,8,9]])
>>> b = np.array([[False,True,False],[True,False,False],[False,False,True]])
>>> a[b]
array([2, 4, 9])

Vous pouvez également ajouter un arange et faire une sélection directe à ce sujet, bien que cela dépende de la façon dont vous générez votre tableau booléen et de l'apparence de votre code YMMV.

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

J'espère que cela vous aidera, faites-moi savoir si vous avez d'autres questions.

73
Slater Victoroff

Vous pouvez faire quelque chose comme ça:

In [7]: a = np.array([[1, 2, 3],
   ...: [4, 5, 6],
   ...: [7, 8, 9]])

In [8]: lst = [1, 0, 2]

In [9]: a[np.arange(len(a)), lst]
Out[9]: array([2, 4, 9])

Plus d'informations sur l'indexation des tableaux multidimensionnels: http://docs.scipy.org/doc/numpy/user/basics.indexing.html#indexing-multi-dimensional-arrays

25
Ashwini Chaudhary

Un moyen simple pourrait ressembler à ceci:

In [1]: a = np.array([[1, 2, 3],
   ...: [4, 5, 6],
   ...: [7, 8, 9]])

In [2]: y = [1, 0, 2]  #list of indices we want to select from matrix 'a'

range(a.shape[0]) renverra array([0, 1, 2])

In [3]: a[range(a.shape[0]), y] #we're selecting y indices from every row
Out[3]: array([2, 4, 9])
3
Dhaval Mayatra

Vous pouvez le faire en utilisant l'itérateur. Comme ça:

np.fromiter((row[index] for row, index in Zip(X, Y)), dtype=int)

Temps:

N = 1000
X = np.zeros(shape=(N, N))
Y = np.arange(N)

#@Aशwini चhaudhary
%timeit X[np.arange(len(X)), Y]
10000 loops, best of 3: 30.7 us per loop

#mine
%timeit np.fromiter((row[index] for row, index in Zip(X, Y)), dtype=int)
1000 loops, best of 3: 1.15 ms per loop

#mine
%timeit np.diag(X.T[Y])
10 loops, best of 3: 20.8 ms per loop
2
Kei Minagawa

En tant qu'approche numpythonique pure, vous pouvez obtenir les indices du deuxième axe en utilisant np.take(), puis la diagonale du résultat serait votre sortie attendue:

np.diagonal(np.take(arr, idx, axis=1))

Démo:

>>> arr = np.array([[1, 2, 3],
... [4, 5, 6],
... [7, 8, 9]])
>>> 
>>> idx = [1, 0, 2]
>>> 
>>> np.diagonal(np.take(arr, idx, axis=1))
array([2, 4, 9])
1
Kasrâmvd

Une autre façon intelligente consiste à transposer d'abord le tableau et à l'indexer par la suite. Enfin, prenez la diagonale, c'est toujours la bonne réponse.

X = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]])
Y = np.array([1, 0, 2, 2])

np.diag(X.T[Y])

étape par étape:

Tableaux originaux:

>>> X
array([[ 1,  2,  3],
       [ 4,  5,  6],
       [ 7,  8,  9],
       [10, 11, 12]])

>>> Y
array([1, 0, 2, 2])

Transposer pour permettre de bien l'indexer.

>>> X.T
array([[ 1,  4,  7, 10],
       [ 2,  5,  8, 11],
       [ 3,  6,  9, 12]])

Obtenez les lignes dans l'ordre Y.

>>> X.T[Y]
array([[ 2,  5,  8, 11],
       [ 1,  4,  7, 10],
       [ 3,  6,  9, 12],
       [ 3,  6,  9, 12]])

La diagonale devrait maintenant devenir claire.

>>> np.diag(X.T[Y])
array([ 2,  4,  9, 12]
0
Thomas Devoogdt