# 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