tutorial - transpuesta de una matriz en python numpy

usando una matriz numpy como índices del segundo dim de otra matriz?

Question

Esta pregunta ya tiene una respuesta aquí:

• Indexación de una matriz por otra en numpy 3 respuestas

Por ejemplo, tengo dos matrices numpy,

A = np.array(
  [[0,1], 
   [2,3], 
   [4,5]])
B = np.array(
  [[1],
   [0],
   [1]], dtype=''int'')

y quiero extraer un elemento de cada fila de A , y ese elemento está indexado por B , así que quiero los siguientes resultados:

C = np.array(
  [[1],
   [2],
   [5]])

Intenté A[:, B.ravel()] , pero emitirá B , no lo que quiero. También examiné np.take , parece que no es la solución correcta para mi problema.

Sin embargo, podría usar np.choose transponiendo A ,

np.choose(B.ravel(), A.T)

pero alguna otra mejor solución?

Answer 1

Puede usar NumPy''s purely integer array indexing :

A[np.arange(A.shape[0]),B.ravel()]

Ejecución de muestra:

In [57]: A
Out[57]: 
array([[0, 1],
       [2, 3],
       [4, 5]])

In [58]: B
Out[58]: 
array([[1],
       [0],
       [1]])

In [59]: A[np.arange(A.shape[0]),B.ravel()]
Out[59]: array([1, 2, 5])

Tenga en cuenta que si B es una matriz 1D o una lista de dichos índices de columna, simplemente puede omitir la operación de aplanamiento con .ravel() .

Ejecución de muestra:

In [186]: A
Out[186]: 
array([[0, 1],
       [2, 3],
       [4, 5]])

In [187]: B
Out[187]: [1, 0, 1]

In [188]: A[np.arange(A.shape[0]),B]
Out[188]: array([1, 2, 5])

Answer 2

C = np.array([A[i][j] for i,j in enumerate(B)])