5

I have the following 2D array 

A=([[1, 2, 3, 4],
    [5, 6, 7, 8],
    [9, 10, 11, 12],
    [13, 14, 15, 16])

And I want to replace the main diagonal by the array 

a = ([0,2,15,20])

Therefore, the results must be

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

I tried with np.diag(a, k=0) but it doesn't work because np.diag() creates a diagonal 2D array with the array "a".

Is there a way to do that with numpy? The above example is the simplest one. I would like to be able to change not only the mail diagonal but all diagonals.

DiegoDZ
  • 189
  • 1
  • 11

2 Answers2

5

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])
>>> a
array([[ 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)
>>> A
array([[  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.

Willem Van Onsem
  • 443,496
  • 30
  • 428
  • 555
  • 1
    I love your answer (discovering numpy functions is lovable), but he asked for the main diagonal, as long as **any** diagonals if I am not mistaken. – IMCoins Feb 01 '18 at 18:24
  • Thanks, it is a good answer. I have changed the array "a" because it was a bad example with all the elements equal to 20. – DiegoDZ Feb 01 '18 at 18:25
  • @IMCoins: well this fills only the *main* diagonal (it is even in the documentation) :) – Willem Van Onsem Feb 01 '18 at 18:26
  • @DiegoDZ: that is not the problem, you can also pass an iterable, and it will then *not* perform broadcasting. – Willem Van Onsem Feb 01 '18 at 18:28
3

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]
Alex
  • 18,484
  • 8
  • 60
  • 80