changing the value of a numpy array based on an offset

Viewed 48

I have a 2 dimensional numpy array

For each column, I want to add 1 to the Nth row where N is the value in the zeroth row of that column.

How can I do this?

A=np.zeros(100)
A=np.reshape(A,[20,5])
A[0]=[5,2,4,1,3]

I want to add 1 to A[5,0], A[2,1], A[4,2], A[1,3] and A[3,4]

1 Answers

Following the following logic: add 1 to A[5,0], A[2,1], A[4,2], A[1,3] and A[3,4], you can do so using advanced indexing:

indices = [5,2,4,1,3]
A[indices, np.arange(len(indices))] = 1

print(A)
array([[0., 0., 0., 0., 0.],
       [0., 0., 0., 1., 0.],
       [0., 1., 0., 0., 0.],
       [0., 0., 0., 0., 1.],
       [0., 0., 1., 0., 0.],
       [1., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0.],
       ...
Related