numpy: how to get a max from an argmax result

Question

I have a numpy array of arbitrary shape, e.g.:

a = array([[[ 1, 2], [ 3, 4], [ 8, 6]], [[ 7, 8], [ 9, 8], [ 3, 12]]]) a.shape = (2, 3, 2)

and a result of argmax over the last axis:

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

I'd like to get max:

np.max(a, axis=-1) = array([[ 2, 4, 8], [ 8, 9, 12]])

But without recalculating everything. I've tried:

a[np.arange(len(a)), np.argmax(a, axis=-1)]

But got:

IndexError: shape mismatch: indexing arrays could not be broadcast together with shapes (2,) (2,3)

How to do it? Similar question for 2-d: numpy 2d array max/argmax

What is reshaped_x?

dnalow
– dnalow

2016-11-01 09:41:44 +00:00
Commented Nov 1, 2016 at 9:41 — dnalow
– dnalow, Commented Nov 1, 2016 at 9:41
Sorry, it should be a. Correcting now.

sygi
– sygi

2016-11-01 09:42:16 +00:00
Commented Nov 1, 2016 at 9:42 — sygi
– sygi, Commented Nov 1, 2016 at 9:42

Community · Accepted Answer · 2017-05-23 11:48:18Z

11

You can use advanced indexing -

In [17]: a Out[17]: array([[[ 1, 2], [ 3, 4], [ 8, 6]], [[ 7, 8], [ 9, 8], [ 3, 12]]]) In [18]: idx = a.argmax(axis=-1) In [19]: m,n = a.shape[:2] In [20]: a[np.arange(m)[:,None],np.arange(n),idx] Out[20]: array([[ 2, 4, 8], [ 8, 9, 12]])

For a generic ndarray case of any number of dimensions, as stated in the comments by @hpaulj, we could use np.ix_, like so -

shp = np.array(a.shape) dim_idx = list(np.ix_(*[np.arange(i) for i in shp[:-1]])) dim_idx.append(idx) out = a[dim_idx]

edited May 23, 2017 at 11:48

CommunityBot

11 silver badge

answered Nov 1, 2016 at 9:56

Divakar

222k19 gold badges273 silver badges374 bronze badges

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

6 Comments

sygi Over a year ago

Ok, but how to do it with an array of arbitrary shape and not exactly 3 dimensions?

hpaulj Over a year ago

Generate the first dimensions with np.xi_ and tuple concatenate idx.

Divakar Over a year ago

@hpaulj Thanks, that should do it! Updated the post with that.

sygi Over a year ago

Thanks. Is there a typo somewhere? It seems dims are not used.

Divakar Over a year ago

@sygi Sorry, was trying out few stuffs and in the end that wasn't required. Removed that line.

|

Nopelm · Accepted Answer · 2024-12-12 13:02:43Z

For arbitrary dimensions and an arbitrary axis on which to operate, the following solution based on numpy.take_along_axis can be used:

import numpy as np # define array as in question a = np.array([ [[ 1, 2], [ 3, 4], [ 8, 6]], [[ 7, 8], [ 9, 8], [ 3, 12]]] ) # define axis of operation ax = -1 # apply argmax using the keepdims option max_indices = np.argmax(a, axis=ax, keepdims=True) # select elements corresponding to indices max_elems = np.take_along_axis(a, max_indices, axis=ax) # squeeze to remove the singleton dimension corresponding to the axis of operation max_elems = max_elems.squeeze(axis=ax)

DavidG · Accepted Answer · 2019-07-13 17:23:33Z

For ndarray with arbitrary shape, you can flatten the argmax indices, then recover the correct shape, as so:

idx = np.argmax(a, axis=-1) flat_idx = np.arange(a.size, step=a.shape[-1]) + idx.ravel() maximum = a.ravel()[flat_idx].reshape(*a.shape[:-1])

lianwei82cn · Accepted Answer · 2020-04-26 07:41:53Z

For arbitrary-shape arrays, the following should work :)

a = np.arange(5 * 4 * 3).reshape((5,4,3)) # for last axis argmax = a.argmax(axis=-1) a[tuple(np.indices(a.shape[:-1])) + (argmax,)] # for other axis (eg. axis=1) argmax = a.argmax(axis=1) idx = list(np.indices(a.shape[:1]+a.shape[2:])) idx[1:1] = [argmax] a[tuple(idx)]

or

a = np.arange(5 * 4 * 3).reshape((5,4,3)) argmax = a.argmax(axis=0) np.choose(argmax, np.moveaxis(a, 0, 0)) argmax = a.argmax(axis=1) np.choose(argmax, np.moveaxis(a, 1, 0)) argmax = a.argmax(axis=2) np.choose(argmax, np.moveaxis(a, 2, 0)) argmax = a.argmax(axis=-1) np.choose(argmax, np.moveaxis(a, -1, 0))

Collectives™ on Stack Overflow

numpy: how to get a max from an argmax result

4 Answers 4

6 Comments

Comments

Comments

Comments

Linked

Hot Network Questions

Collectives™ on Stack Overflow

4 Answers 4

6 Comments

Comments

Comments

Comments

Linked

Related