changing the values of the diagonal of a matrix in numpy

Question

how can I change the values of the diagonal of a matrix in numpy?

I checked Numpy modify ndarray diagonal, but the function there is not implemented in numpy v 1.3.0

Answer 1

If you're using a version of numpy that doesn't have fill_diagonal (the right way to set the diagonal to a constant) or diag_indices_from, you can do this pretty easily with array slicing:

# assuming a 2d square array
n = mat.shape[0]
mat[range(n), range(n)] = 0

This is much faster than an explicit loop in Python, because the looping happens in C and is potentially vectorized.

One nice thing about this is that you can also fill a diagonal with a list of elements, rather than a constant value (like diagflat, but for modifying an existing matrix rather than making a new one). For example, this will set the diagonal of your matrix to 0, 1, 2, ...:

# again assuming 2d square array
n = mat.shape[0]
mat[range(n), range(n)] = range(n)

If you need to support more array shapes, this is more complicated (which is why fill_diagonal is nice...):

m[list(zip(*map(range, m.shape)))] = 0

(The list call is only necessary in Python 3, where zip returns an iterator.)