Numpy: change max in each row to 1, all other numbers to 0

Question

I'm trying to implement a numpy function that replaces the max in each row of a 2D array with 1, and all other numbers with zero:

>>> a = np.array([[0, 1], ...               [2, 3], ...               [4, 5], ...               [6, 7], ...               [9, 8]]) >>> b = some_function(a) >>> b [[0. 1.]  [0. 1.]  [0. 1.]  [0. 1.]  [1. 0.]] 

What I've tried so far

def some_function(x):     a = np.zeros(x.shape)     a[:,np.argmax(x, axis=1)] = 1     return a  >>> b = some_function(a) >>> b [[1. 1.]  [1. 1.]  [1. 1.]  [1. 1.]  [1. 1.]] 

Answer 1

Method #1, tweaking yours:

>>> a = np.array([[0, 1], [2, 3], [4, 5], [6, 7], [9, 8]]) >>> b = np.zeros_like(a) >>> b[np.arange(len(a)), a.argmax(1)] = 1 >>> b array([[0, 1],        [0, 1],        [0, 1],        [0, 1],        [1, 0]]) 

[Actually, range will work just fine; I wrote arange out of habit.]

Method #2, using max instead of argmax to handle the case where multiple elements reach the maximum value:

>>> a = np.array([[0, 1], [2, 2], [4, 3]]) >>> (a == a.max(axis=1)[:,None]).astype(int) array([[0, 1],        [1, 1],        [1, 0]]) 

Answer 2

I prefer using numpy.where like so:

a[np.where(a==np.max(a))] = 1