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.

Divakar · Accepted Answer · 2017-12-06 16:21:25Z

20

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

edited Dec 6, 2017 at 16:21

answered Dec 6, 2017 at 16:00

Divakar

222k19 gold badges273 silver badges374 bronze badges

Sign up to request clarification or add additional context in comments.

2 Comments

Paul Panzer Over a year ago

I like the second one, ( np.equal.outer(a.argmax(axis=1), range(a.shape[1])))

feedMe Over a year ago

I ended up using a variation of Method #2. It also works when the axis=1 arrays contain more than one item corresponding to the max value. argmax returns the first occurrence as True only, which is the behavior that works best for my problem

Paul Panzer · Accepted Answer · 2017-12-06 15:58:40Z

3

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.

edited Dec 6, 2017 at 15:58

answered Dec 6, 2017 at 15:52

Paul Panzer

53.3k3 gold badges60 silver badges103 bronze badges

2 Comments

Divakar Over a year ago

np.eye was one shorter alternative, but needs dtype= there.

Paul Panzer Over a year ago

@Divakar (a.max(1)==a.T).T looks hard to beat in that respect.

kmario23 · Accepted Answer · 2017-12-06 16:06:38Z

2

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

edited Dec 6, 2017 at 16:06

answered Dec 6, 2017 at 15:54

kmario23

62.1k17 gold badges174 silver badges160 bronze badges

Comments

B. M. · Accepted Answer · 2017-12-06 15:56:39Z

0

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.

answered Dec 6, 2017 at 15:56

B. M.

18.7k2 gold badges40 silver badges56 bronze badges

Comments

VladVin · Accepted Answer · 2019-11-26 08:37:44Z

0

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

answered Nov 26, 2019 at 8:37

VladVin

6437 silver badges13 bronze badges

Collectives™ on Stack Overflow

Mask from max values in numpy array, specific axis

5 Answers 5

2 Comments

2 Comments

Comments

Comments

Comments

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

5 Answers 5

2 Comments

2 Comments

Comments

Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Linked

Related