sygi - 1 year ago 208
Python Question

# numpy: how to get a max from an argmax result

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

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]
``````
Recommended from our users: Dynamic Network Monitoring from WhatsUp Gold from IPSwitch. Free Download