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

[Ken Basye]

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

I tried several approaches with take() but now that I understand how 
take() works when you give it an axis argument it seems like this isn't 
going to do what I want.  Still, it seems like there should be some 
shortcut...

    TIA,
       Ken