1 in each row of an all zeros numpy array at particular column index given by another numpy array

Question

I have an array of some given size eg 4x4 with all zeros,

a = np.zeros((4,4))

and I want to put 1 in each row at the column index given by another array

b = np.array([0,1,2,1])

so the resulted array should look like this,

a = 
1   0   0   0
0   1   0   0
0   0   1   0
0   1   0   0

How can I do this for a large array of size (mxn) given b of size (mx1) .

Thank You and Best Regards,

Answer 1

You can use this simple way of indexing 2D array:

>>> a[np.arange(len(a)), b] = 1
>>> a
array([[1., 0., 0., 0.],
       [0., 1., 0., 0.],
       [0., 0., 1., 0.],
       [0., 1., 0., 0.]])

Answer 2

Loop through rows in a and values in b using zip() :

for row, idx in zip(a, b):
    row[idx] = 1