NumPy: Find the maximum value from each row set it to 1 and rest as 0

Question

I have a 2D numpy array as this,

array([[ 0.49596769,  1.15846407, -1.38944733],
       [-0.47042814, -0.07512128 , 1.90417981]], dtype=float32)

I want to find the maximum for each row and change it to 1 and rest as 0. Like this.

array([[ 0.,  1.,  0.],
       [ 0.,  0.,  1.]], dtype=float32)

What is the most efficient way to get it done using numpy?

Answer 1

One approach would be -

(a == a.max(axis=1, keepdims=1)).astype(float)

Sample run -

In [43]: a
Out[43]: 
array([[ 0.49596769,  1.15846407, -1.38944733],
       [-0.47042814, -0.07512128,  1.90417981]])

In [44]: (a == a.max(axis=1, keepdims=1)).astype(float)
Out[44]: 
array([[ 0.,  1.,  0.],
       [ 0.,  0.,  1.]])

---

If there are multiple ones in a row with the same maximum value and you would like to set only the first one as 1 -

idx = a.argmax(axis=1)
out = (idx[:,None] == np.arange(a.shape[1])).astype(float)

Sample run -

In [49]: a
Out[49]: 
array([[2, 4, 4],
       [3, 4, 5]])

In [50]: (a == a.max(axis=1, keepdims=1)).astype(float)
Out[50]: 
array([[ 0.,  1.,  1.],
       [ 0.,  0.,  1.]])

In [51]: idx = a.argmax(axis=1)

In [52]: (idx[:,None] == np.arange(a.shape[1])).astype(float)
Out[52]: 
array([[ 0.,  1.,  0.],
       [ 0.,  0.,  1.]])

For performance, we can have an initialization based approach -

def initialization_based(a): 
    idx = a.argmax(axis=1)
    out = np.zeros_like(a,dtype=float)
    out[np.arange(a.shape[0]), idx] = 1
    return out