CODEWITHSUNDEEP
pythonnumpy
feedMe
feedMe

Reputation: 3727

Mask from max values in numpy array, specific axis

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.

Upvotes: 11

Views: 11514

Answers (5)

VladVin
VladVin

Reputation: 643

In a multi-dimensional case you can also use np.indices. Let's suppose you have an array:

a = np.array([[
    [0, 1, 2],
    [3, 8, 5],
    [6, 7, -1],
    [9, 5, 8]],[
    [5, 2, 8],
    [7, 6, -3],
    [-1, 2, 1],
    [3, 5, 6]]
])

you can access argmax values calculated for axis 0 like so:

k = np.zeros((2, 4, 3), np.bool)
k[a.argmax(0), ind[0], ind[1]] = 1

The output would be:

array([[[False, False, False],
        [False,  True,  True],
        [ True,  True, False],
        [ True,  True,  True]],

       [[ True,  True,  True],
        [ True, False, False],
        [False, False,  True],
        [False, False, False]]])

Upvotes: 0

Divakar
Divakar

Reputation: 221664

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

Upvotes: 20

kmario23
kmario23

Reputation: 61445

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

Upvotes: 2

Paul Panzer
Paul Panzer

Reputation: 53089

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.

Upvotes: 3

B. M.
B. M.

Reputation: 18668

in on line : np.equal(a.max(1)[:,None],a) or np.equal(a.max(1),a.T).T .

But this can lead to several ones in a row.

Upvotes: 0

Related Questions