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


93

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?


A resposta é melhor para você? stackoverflow.com/a/17081678/5046896
GoingMyWay

Respostas:


104

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.


13
1 para o exemplo usando arange. Isso foi particularmente útil para mim para recuperar blocos diferentes de matrizes múltiplas (basicamente o caso 3D deste exemplo)
Griddo

1
Olá, você poderia explicar por que temos que usar em arangevez de :? Sei que o seu jeito funciona e o meu não, mas gostaria de entender por quê.
marcotama

@tamzord porque é um array numpy e não uma lista python vanilla, então a :sintaxe não funciona da mesma maneira.
Slater Victoroff

1
@SlaterTyranus, obrigado por responder. Meu entendimento, após algumas leituras, é que misturar :com indexação avançada significa: "para cada subespaço ao longo :, aplique a indexação avançada fornecida". Meu entendimento está correto?
marcotama

@tamzord explica o que você entende por "
subespaço

36

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


2
lutando para entender por que o arange é necessário em vez de simplesmente ':' ou range.
MadmanLee

@MadmanLee Oi, usando :produzirá várias len(a)vezes dos resultados, em vez disso, indicando que o índice de cada linha imprimirá os resultados antecipados.
GoingMyWay

1
Acho que essa é exatamente a maneira certa e elegante de resolver esse problema.
GoingMyWay

6

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

1
Por favor, considere adicionar explicações.
souki

@souki Eu adicionei uma explicação agora. Obrigado
Dhaval Mayatra

6

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.


3

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

1
O OP mencionou que ele deve ser executado rapidamente em grandes arrays, então seus benchmarks não são muito representativos. Estou curioso para saber como é o desempenho do seu último método para matrizes (muito) maiores!

@moarningsun: Atualizado. np.diag(X.T[Y])é tão lento ... Mas np.diag(X.T)é tão rápido (10us). Não sei por quê.
Kei Minagawa

0

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]

1
Isso funciona tecnicamente e parece muito elegante. No entanto, acho que essa abordagem explode completamente quando você está lidando com matrizes grandes. No meu caso, o NumPy engoliu 30 GB de swap e preencheu meu SSD. Eu recomendo usar a abordagem de indexação avançada.
nefasto
Ao utilizar nosso site, você reconhece que leu e compreendeu nossa Política de Cookies e nossa Política de Privacidade.
Licensed under cc by-sa 3.0 with attribution required.