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 e.g 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,

mathfux · Accepted Answer · 2020-08-20 13:32:41Z

3

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.]])

answered Aug 20, 2020 at 13:32

mathfux

5,9792 gold badges21 silver badges38 bronze badges

Sign up to request clarification or add additional context in comments.

Comments

Michał Urbanowicz · Accepted Answer · 2020-08-20 13:28:12Z

0

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

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

answered Aug 20, 2020 at 13:28

717 bronze badges

Thank you for your answer, but can it be achieved without for loop in a more efficient manner?