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

numpy, python

Solution

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]])

Problem

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.]] ```

Original source