Assign value based on index from another array numpy

Question

I have a indices array like this:

idx = np.array([3,4,1], [0,0,0], [1,4,1], [2,0,2]]

And an array of zeros A with shape 4x5

I would like to make all the indices in idx of A to be 1

For the above example, the final array should be:

[[0,1,0,1,1],  # values at index 3,4,1 are 1
 [1,0,0,0,0],  # value at index 0 is 1
 [0,1,0,0,1],  # values at index 1,4 are 1
 [1,0,1,0,0]]  # values at index 0,2 are 1

How can this be done in numpy?

Answer 1

Use fancy indexing:

A = np.zeros((4, 5), dtype=int)

A[np.arange(len(idx))[:,None], idx] = 1

Or numpy.put_along_axis:

A = np.zeros((4, 5), dtype=int)

np.put_along_axis(A, idx, 1, axis=1)

updated A:

array([[0, 1, 0, 1, 1],
       [1, 0, 0, 0, 0],
       [0, 1, 0, 0, 1],
       [1, 0, 1, 0, 0]])