Can I improve my Numpy solution to an exercise? [closed]

Question

I have been asked to use the following set of column indices: y = np.array([3, 0, 4, 1])

to turn into 1 all the elements in the following matrix: x = np.zeros(shape = (4, 5)) that have y as starting column and rows given by the position of y. Just to be clear.

The final result has to be the following:

[[0. 0. 0. 1. 1.]      
 [1. 1. 1. 1. 1.]  
 [0. 0. 0. 0. 1.]  
 [0. 1. 1. 1. 1.]]

For example:

y[0] = 3,

then row 0, columns 3 and 4 need to be equal to 1.

I did it like this:

for (idx, num) in enumerate(y): x[idx, num:] = 1

Can this result be written differently and/or improved by using other Numpy functions (for example, using vectorization)?

Answer 1

Lots of ways of doing this. For example, since x is basically a boolean mask, you can compute a mask and turn it into whatever type you want:

x = (np.arange(5) < y[:, None]).astype(float)

You might also use np.where to avoid the conversion:

x = np.where(np.arange(5) < y[:, None], 1.0, 0.0)

Answer 2

You can try np.indices like below

(np.indices(x.shape)[1]>=y[:,None]).astype(int)

which gives

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