[Numpy-discussion] Extracting values from one array corresponding to argmax elements in another array

[Keith Goodman]

On Mon, Apr 5, 2010 at 8:44 AM, Ken Basye <kbasye1@jhu.edu> wrote:
> Hi Folks,
>  I have two arrays, A and B, with the same shape.  I want to find the
> highest values in A along some axis, then extract the corresponding
> values from B.  I can get the highest values in A with A.max(axis=0) and
> the indices of these highest values with A.argmax(axis=0).  I'm trying
> to figure out a loop-free way to extract the corresponding elements from
> B using these indices.  Here's code with a loop that will do what I want
> for two-dimensional arrays:
>
>  >>> a
> array([[ 100.,    0.,    0.],
>       [   0.,  100.,  100.],
>       [   0.,    0.,    0.]])
>
>  >>> a.max(axis=0)
> array([ 100.,  100.,  100.])
>
>  >>> sel = a.argmax(axis=0)
>  >>>sel
> array([0, 1, 1])
>
>  >>> b = np.arange(9).reshape((3,3))
>  >>> b
> array([[0, 1, 2],
>       [3, 4, 5],
>       [6, 7, 8]])
>
>  >>> b_best = np.empty(3)
>  >>> for i in xrange(3):
> ...    b_best[i] = b[sel[i], i]
> ...
>  >>> b_best
> array([ 0.,  4.,  5.])

Here's one way:

>> b[a.argmax(axis=0), range(3)]
   array([0, 4, 5])