setting values in numpy matrix with array of indices

Question

Consider a mxn matrix A. I computed the indices of the maximum of each row, yielding an array of dimension m.

How can I use this array of indices to set the values in a second matrix B of the same shape as A in each row to 0.

Example:

     A = [[1,2] ,[3,4]] 

     np.argmax(A,1) --> [1,1]

     B = [[1,1] ,[1,1]]

I want to have:

  B = [[1 0] ,[1 0]]

How can I do that?

warped · Accepted Answer · 2022-01-16 18:47:09Z

1

you can use:

B[np.arange(B.shape[0]),np.argmax(A,1)] = 0

answered Jan 16, 2022 at 18:33

warped

9,6655 gold badges26 silver badges55 bronze badges

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

2 Comments

If I use B= np.ones((20, 4), dtype='bool'), then this does not work? A is also of shape 20x4

@Sarah try the edit