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

Sorry, it should be `a`. Correcting now. – sygi Nov 01 '16 at 09:42 — sygi, Nov 01 '16 at 09:42

score 11 · Accepted Answer · edited May 23 '17 at 11:48

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 '17 at 11:48

Community

1
1

answered Nov 01 '16 at 09:56

Divakar

218,885
19
262
358

1

Ok, but how to do it with an array of arbitrary shape and not exactly 3 dimensions? – sygi Nov 01 '16 at 10:40
1

Generate the first dimensions with `np.xi_` and tuple concatenate `idx`. – hpaulj Nov 01 '16 at 11:48
@hpaulj Thanks, that should do it! Updated the post with that. – Divakar Nov 01 '16 at 13:45
Thanks. Is there a typo somewhere? It seems `dims` are not used. – sygi Nov 01 '16 at 16:21
1

@sygi Sorry, was trying out few stuffs and in the end that wasn't required. Removed that line. – Divakar Nov 01 '16 at 16:37
The FutureWarning says we want tuple `a[tuple([*np.ix_(*[np.arange(i) for i in a.shape[:-1]]), idxs])]` – Franz Knülle Apr 01 '23 at 11:23

score 0 · Answer 2 · answered Jul 13 '19 at 17:23

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 · Answer 3 · 2020-04-26T07:41:53.520

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

numpy: how to get a max from an argmax result

3 Answers3

Linked

Related