Return the diagonal of the matrix product A @ B without computing the full (n, n) product first.
A @ B
(n, n)
Input
A = [[0. 1. 2.] [3. 4. 5.] [6. 7. 8.]] B = [[ 9. 10. 11.] [12. 13. 14.] [15. 16. 17.]]
Output
[ 42. 162. 300.]
This one's part of Premium. Unlock the full NumPy track plus every other premium problem on the site.
solve(...)