Numpy: Assignment and Indexing as Matlab

Question

Sometimes is useful to assign arrays with one index only. In Matlab this is straightforward:

M = zeros(4);
M(1:5:end) = 1
M =

   1   0   0   0
   0   1   0          


        

Answer 1

You could do this using list indices:

M = np.zeros((4,4))
M[range(4), range(4)] = 1
print M
# [[ 1.  0.  0.  0.]
#  [ 0.  1.  0.  0.]
#  [ 0.  0.  1.  0.]
#  [ 0.  0.  0.  1.]]

In this case you could also use np.identity(4)