# Slice of 2d numpy array with another array

## Question:

I have a quite large 2d array, and I need to get both the index of the maximum value in axis 1, and the maximum value itself. I can retrieve these two values as follows:

``````import numpy as np
a = np.arange(27).reshape(9, 3)
idx = np.argmax(a, axis=1)
max_val = np.max(a, axis=1)
``````

However, since I have already found the index of the maximum value, it feels like I should be able to construct the array of maximum values using idx without having to look up the value again.

I realise I can use `np.choose(idx, a.T)` but this involves transposing the matrix which will be much more expensive than just using `max`. I can do something like `np.array([a[i][idx[i]] for i in range(len(a))])` but this involves creating a list which again seems more expensive that just calling `np.max`.

Is there any way to slice `a` with `idx` in numpy without restructuring the array?

Your `a` and `argmax`:

``````In [602]: a
Out[602]:
array([[ 0,  1,  2],
[ 3,  4,  5],
[ 6,  7,  8],
[ 9, 10, 11],
[12, 13, 14],
[15, 16, 17],
[18, 19, 20],
[21, 22, 23],
[24, 25, 26]])

In [603]: idx
Out[603]: array([2, 2, 2, 2, 2, 2, 2, 2, 2], dtype=int64)
``````

A common way of using that index array:

``````In [606]: a[np.arange(a.shape[0]),idx]
Out[606]: array([ 2,  5,  8, 11, 14, 17, 20, 23, 26])
``````

A newer tool, that may be easier to use (if not familiar with the first):

``````In [607]: np.take_along_axis(a,idx[:,None],1)
Out[607]:
array([[ 2],
[ 5],
[ 8],
[11],
[14],
[17],
[20],
[23],
[26]])
``````
Categories: questions Tags: , , ,
Answers are sorted by their score. The answer accepted by the question owner as the best is marked with
at the top-right corner.