Mask from max values in numpy array, specific axis

Create an identity matrix and select from its rows using argmax on your array:

np.identity(a.shape[1], bool)[a.argmax(axis=1)]
# array([[False,  True],
#        [ True, False],
#        [False,  True]], dtype=bool)

Please note that this ignores ties, it just goes with the value returned by argmax.


Method #1

Using broadcasting, we can use comparison against the max values, while keeping dims to facilitate broadcasting -

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

Sample run -

In [83]: a
Out[83]: 
array([[0, 1],
       [2, 1],
       [4, 8]])

In [84]: a.max(axis=1,keepdims=1) == a
Out[84]: 
array([[False,  True],
       [ True, False],
       [False,  True]], dtype=bool)

Method #2

Alternatively with argmax indices for one more case of broadcasted-comparison against the range of indices along the columns -

In [92]: a.argmax(axis=1)[:,None] == range(a.shape[1])
Out[92]: 
array([[False,  True],
       [ True, False],
       [False,  True]], dtype=bool)

Method #3

To finish off the set, and if we are looking for performance, use intialization and then advanced-indexing -

out = np.zeros(a.shape, dtype=bool)
out[np.arange(len(a)), a.argmax(axis=1)] = 1

You're already halfway in the answer. Once you compute the max along an axis, you can compare it with the input array and you'll have the required binary mask!

In [7]: maxx = np.amax(a, axis=1)

In [8]: maxx
Out[8]: array([1, 2, 8])

In [12]: a >= maxx[:, None]
Out[12]: 
array([[False,  True],
       [ True, False],
       [False,  True]], dtype=bool)

Note: This uses NumPy broadcasting when doing the comparison between a and maxx

Tags:

Python

Numpy