NumPy selecionando um índice de coluna específico por linha usando uma lista de índices

Estou lutando para selecionar as colunas específicas por linha de uma matriz NumPy.

Suponha que eu tenha a seguinte matriz que eu chamaria de X:

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

Eu também tenho um listíndice de coluna para cada linha que eu chamaria de Y:

[1, 0, 2]

Eu preciso obter os valores:

[2]
[4]
[9]

Em vez de um listcom índices Y, também posso produzir uma matriz com a mesma forma em Xque cada coluna é um bool/ intno intervalo de 0-1 valor, indicando se esta é a coluna necessária.

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

Sei que isso pode ser feito iterando a matriz e selecionando os valores de coluna de que preciso. No entanto, isso será executado com frequência em grandes arrays de dados e é por isso que deve ser executado o mais rápido possível.

Eu queria saber se existe uma solução melhor?

Se você tem uma matriz booleana, pode fazer a seleção direta com base nisso:

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

Para acompanhar seu exemplo inicial, você pode fazer o seguinte:

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

Você também pode adicionar um arangee fazer uma seleção direta nisso, embora dependendo de como você está gerando sua matriz booleana e como seu código se parece com 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])

Espero que isso ajude, entre em contato se tiver mais perguntas.

Você pode fazer algo assim:

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

Mais sobre indexação de matrizes multidimensionais: http://docs.scipy.org/doc/numpy/user/basics.indexing.html#indexing-multi-dimensional-arrays

Uma maneira simples pode ser assim:

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

numpyVersões recentes adicionaram um take_along_axis(e put_along_axis) que faz essa indexação de forma limpa.

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

Funciona da mesma forma que:

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

mas com manuseio de eixo diferente. Destina-se especialmente a aplicar os resultados de argsorte argmax.

Você pode fazer isso usando o iterador. Como isso:

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

Outra maneira inteligente é primeiro transpor a matriz e indexá-la depois disso. Finalmente, pegue a diagonal, é sempre a resposta certa.

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 a passo:

Matrizes originais:

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

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

Transponha para que seja possível indexá-lo corretamente.

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

Obtenha as linhas na ordem Y.

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

A diagonal agora deve ficar clara.

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