Find the maximum values of a matrix in rows axis and replace other values to zero-CodePudding

A = [[2,2,4,2,2,2]
     [2,6,2,2,2,2]
     [2,2,2,2,8,2]]

I want matrix B to be equal to:

B = [[0,0,4,0,0,0]
     [0,6,0,0,0,0]
     [0,0,0,0,8,0]]

So I want to find the maximum value of each row and replace other values with 0. Is there any way to do this without using for loops? Thanks in advance for your comments.

CodePudding user response：

Instead of looking at the argmax, you could take the max values for each row directly, then mask the elements which are lower and replace them with zeros:

Inplace this would look like (here True stands for keepdims=True):

>>> A[A < A.max(1, True)] = 0

>>> A
array([[0, 0, 4, 0, 0, 0],
       [0, 6, 0, 0, 0, 0],
       [0, 0, 0, 0, 8, 0]])

An out of place alternative is to use np.where:

>>> np.where(A == A.max(1, True), A, 0)
array([[0, 0, 4, 0, 0, 0],
       [0, 6, 0, 0, 0, 0],
       [0, 0, 0, 0, 8, 0]])