NumPy selezionando l'indice di colonna specifico per riga utilizzando un elenco di indici

Sto lottando per selezionare le colonne specifiche per riga di una matrice NumPy.

Supponiamo di avere la seguente matrice che chiamerei X:

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

Ho anche un listindice di colonna per ogni riga che chiamerei Y:

[1, 0, 2]

Ho bisogno di ottenere i valori:

[2]
[4]
[9]

Invece di una listcon indici Y, posso anche produrre una matrice con la stessa forma di Xdove ogni colonna è un bool/ intnell'intervallo 0-1 valore, indicando se questa è la colonna richiesta.

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

So che questo può essere fatto iterando sull'array e selezionando i valori di colonna di cui ho bisogno. Tuttavia, questo verrà eseguito frequentemente su grandi array di dati ed è per questo che deve essere eseguito il più velocemente possibile.

Mi chiedevo quindi se esiste una soluzione migliore?

Se hai un array booleano puoi eseguire la selezione diretta in base a quello in questo modo:

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

Per andare avanti con il tuo esempio iniziale potresti fare quanto segue:

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

Puoi anche aggiungere arangee fare una selezione diretta su quello, anche se a seconda di come stai generando il tuo array booleano e come appare il tuo codice 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])

Spero che questo ti aiuti, fammi sapere se hai altre domande.

Puoi fare qualcosa del genere:

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

Ulteriori informazioni sull'indicizzazione di array multidimensionali: http://docs.scipy.org/doc/numpy/user/basics.indexing.html#indexing-multi-dimensional-arrays

Un modo semplice potrebbe assomigliare a:

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]) sarà di ritorno 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])

Le numpyversioni recenti hanno aggiunto un take_along_axis(e put_along_axis) che esegue questa indicizzazione in modo pulito.

In [101]: a = np.arange(1,10).reshape(3,3)                                                             
In [102]: b = np.array([1,0,2])                                                                        
In [103]: np.take_along_axis(a, b[:,None], axis=1)                                                     
Out[103]: 
array([[2],
       [4],
       [9]])

Funziona allo stesso modo di:

In [104]: a[np.arange(3), b]                                                                           
Out[104]: array([2, 4, 9])

ma con diversa gestione degli assi. È particolarmente mirato ad applicare i risultati di argsorte argmax.

Puoi farlo usando l'iteratore. Come questo:

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

Tempo:

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

Un altro modo intelligente consiste nel trasporre prima l'array e successivamente indicizzarlo. Infine, prendi la diagonale, è sempre la risposta giusta.

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

Passo dopo passo:

Array originali:

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

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

Trasponi per renderlo possibile indicizzarlo correttamente.

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

Ottieni righe nell'ordine Y.

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

La diagonale dovrebbe ora diventare chiara.

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