I know you can iterate over a 2d matrix using two indexes like this:

import numpy as npA = np.zeros((10,10))for i in range(0,10):for j in range(0,10):if (i==j):A[i,j] = 4

Is there a way of doing this using only one for loop or using slices?

EDIT:

I also need to take into account of when i =/ j, for example:

A = np.zeros((10,10))for i in range(0,10):for j in range(0,10):if (i==j):A[i,j] = 1if (i+1 ==j):A[i,j] = 2if (i-1==j):A[i,j] = 3

You can always collapse multiple loops into one by calculating the components each iteration with the modulo operator like so:

import numpy as npA = np.zeros((10,10))for x in range(100):i = math.floor(x/10)j = x % 10if (i==j):A[i,j] = 1if (i+1 ==j):A[i,j] = 2if (i-1==j):A[i,j] = 3

With only i==j it could be even simpler:

for i in range(10):A[i,i] = 4