Replace diagonals of a 2D array with python [duplicate]
You can use np.fill_diagonal(..)
for that. Like the documentation says:
numpy.fill_diagonal(a, val, wrap=False)
Fill the main diagonal of the given array of any dimensionality.
For example:
np.fill_diagonal(A, 20)
We here thus broadcast 20
over the entire diagonal.
You can also fill the diagonal with different values, like:
np.fill_diagonal(A, [0,2,15,20])
For example:
>>> a = np.zeros((4,4))>>> np.fill_diagonal(a, [0,2,15,20])>>> aarray([[ 0., 0., 0., 0.], [ 0., 2., 0., 0.], [ 0., 0., 15., 0.], [ 0., 0., 0., 20.]])
In case you want to change other diagonals, then it is a matter of mirroring the array. For example for the antidiagonal, we can use:
np.fill_diagonal(A[::-1], -20)
Then we thus get:
>>> A = np.zeros((4,4))>>> np.fill_diagonal(A[::-1], -20)>>> Aarray([[ 0., 0., 0., -20.], [ 0., 0., -20., 0.], [ 0., -20., 0., 0.], [-20., 0., 0., 0.]])
If we do not take superdiagonals and subdiagonals into account, a n dimensional matrix, has nĂ—(n-1) diagonals. We can assign this by mirroring one or more dimensions.
Checkout the numpy docs on indexing into multidimensional arrays.
A[np.arange(A.shape[0]), np.arange(A.shape[1])] = [0,2,15,20]
Note: @WillemVanOnsem's answer is the best answer for filling in main diagonal but this is the best general way for getting/setting any subset of elements in a multidimensional array!
For example to change the other diagonal:
A[-np.arange(1, A.shape[0] + 1), np.arange(A.shape[1])] = [0,2,15,20]