How to replace the diagonal elements of a numpy 2d array with another array?-CodePudding

I have a numpy 2d array:

import numpy as np
A=np.array([[1, 2, 3, 4],
            [5, 6, 7, 8],
            [9, 10, 11, 12],
            [13, 14, 15, 16]])
print (A)

I would like to replace the diagonal elements with a = np.array([0,2,15,20]). That is the desired output should be:

A=[[0, 2, 3, 4],
    [5, 2, 7, 8],
    [9, 10, 15, 12],
    [13, 14, 15, 20]]

I tried with the following code:

import numpy as np
A=np.array([[1, 2, 3, 4],
            [5, 6, 7, 8],
            [9, 10, 11, 12],
            [13, 14, 15, 16]])
a = np.array([0,2,15,20])
print(np.fill_diagonal(A, a))

But it resulted in None

CodePudding user response：

As an alternative method if a be the main array and b the modified values array:

a_mod = a.ravel()
a_mod[::a.shape[0] 1] = b
result = a_mod.reshape(a.shape)

It can handle where the other diagonals of a matrix (instead the main diagonal) is of interest, by some modification. np.fill_diagonal works on the main diagonal.

CodePudding user response：

try this

A[np.arange(4), np.arange(4)] = a

array([[ 0,  2,  3,  4],
       [ 5,  2,  7,  8],
       [ 9, 10, 15, 12],
       [13, 14, 15, 20]])