Diagonal of a Matrix Product
mediumnumpy_linalgImplement diag_of_product(A, B).
Given A of shape (n, k) and B of shape (k, n), return the diagonal of
the product A @ B as a 1D array of length n — i.e. element i is
A[i] . B[:, i]. Do it without forming the full A @ B, and with no
for/while loops.
Constraints
- no scipy, sklearn
- no for/while loops
Your solution
Edit diag_of_product and run the real pytest suite in your browser — no install required. Your code is saved locally.
Loading the Python runtime (first run only)…
Hints
Hint 1
Entry i of diag(A@B) is row i of A dotted with column i of B — np.einsum('ij,ji->i', A, B) does it without building the full product.