(Python) ¿Cómo obtener la diagonal (A * B) sin tener que realizar A * B?

Digamos que tenemos dos matrices A y B y que la matriz C sea A*B (la multiplicación de la matriz no es un elemento). Deseamos obtener solo las entradas diagonales de C , que se pueden realizar a través de np.diagonal(C) . Sin embargo, esto genera una sobrecarga de tiempo innecesaria, porque estamos multiplicando A por B aunque solo necesitamos las multiplicaciones de cada fila en A con la columna de B que tiene el mismo ‘id’, que es la fila 1 de A con la columna 1 de B , fila 2 de A con columna 2 de B y así sucesivamente: las multiplicaciones que forman la diagonal de C ¿Hay alguna manera de lograrlo eficientemente usando Numpy? Quiero evitar el uso de bucles para controlar qué fila se multiplica con qué columna, en su lugar, deseo un método numpy incorporado que realice este tipo de operación para optimizar el rendimiento.

Gracias por adelantado..

Podría usar einsum aquí:

 >>> a = np.random.randint(0, 10, (3,3)) >>> b = np.random.randint(0, 10, (3,3)) >>> a array([[9, 2, 8], [5, 4, 0], [8, 0, 6]]) >>> b array([[5, 5, 0], [3, 5, 5], [9, 4, 3]]) >>> a.dot(b) array([[123, 87, 34], [ 37, 45, 20], [ 94, 64, 18]]) >>> np.diagonal(a.dot(b)) array([123, 45, 18]) >>> np.einsum('ij,ji->i', a,b) array([123, 45, 18]) 

Para matrices más grandes, será mucho más rápido que hacer la multiplicación directamente:

 >>> a = np.random.randint(0, 10, (1000,1000)) >>> b = np.random.randint(0, 10, (1000,1000)) >>> %timeit np.diagonal(a.dot(b)) 1 loops, best of 3: 7.04 s per loop >>> %timeit np.einsum('ij,ji->i', a, b) 100 loops, best of 3: 7.49 ms per loop 

[Nota: originalmente había hecho la versión elementwise, ii,ii->i , en lugar de la multiplicación de matrices. Los mismos trucos de einsum funcionan.]

 def diag(A,B): diags = [] for x in range(len(A)): diags.append(A[x][x] * B[x][x]) return diags 

Creo que el código anterior es lo que estás buscando.