Mask from max values in numpy array, specific axis

Question

Input example:

I have a numpy array, e.g.

a=np.array([[0,1], [2, 1], [4, 8]])

Desired output:

I would like to produce a mask array with the max value along a given axis, in my case axis 1, being True and all others being False. e.g. in this case

mask = np.array([[False, True], [True, False], [False, True]])

Attempt:

I have tried approaches using np.amax but this returns the max values in a flattened list:

>>> np.amax(a, axis=1)
array([1, 2, 8])

and np.argmax similarly returns the indices of the max values along that axis.

>>> np.argmax(a, axis=1)
array([1, 0, 1])

I could iterate over this in some way but once these arrays become bigger I want the solution to remain something native in numpy.

Answer 1

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

Answer 2

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.

Answer 3

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