[FIXED] Replacing diagonal elements with another array in Python

Issue

I have two matrices, A and P. I would like to replace the diagonal elements of A with the elements of P. The desired output is attached.

A=np.array([[0, 1, 1, 0],
       [1, 0, 0, 1],
       [1, 0, 0, 1],
       [0, 1, 1, 0]])

P=np.array([[3, 4],
       [5, 6]])

Desired output:

array([[3, 1, 1, 0],
       [1, 4, 0, 1],
       [1, 0, 5, 1],
       [0, 1, 1, 6]])

Solution

Use fill_diagonal with the flattened P as values for the diagonal:

np.fill_diagonal(A, P.ravel())

NB. the operation is in place

output:

array([[3, 1, 1, 0],
       [1, 4, 0, 1],
       [1, 0, 5, 1],
       [0, 1, 1, 6]])

Answered By - mozway