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?

CodePudding user response：

you can use:

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