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

[Ken Basye]

From: Vincent Schut <schut@sarvision.nl>
> On 04/05/2010 06:06 PM, Keith Goodman wrote:
>   
>> 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])
>>     
>
> Which does not work anymore when your arrays become more-dimensional 
> (like in my case: 4 or more) and the axis you want to select on is not 
> the first/last one. If I recall correctly, I needed to construct the 
> full index arrays for the other dimensions too (like with ogrid I 
> think). So: create the ogrid, replace the one for the dimensions you 
> want the argmax selection to take place on with the argmax parameter, 
> and use those index arrays to index your b array.
> I'd need to look up my source code to be more sure/precise. If anyone 
> would like me to, please let me know. If anyone knows a less elaborate 
> way, also please let us know! :-)
>   
Hi Vincent,
  I'd like to see more about your solution.  For my present purposes, 
Keith's solution was sufficient, but I'm still very interested in a 
solution that's independent of dimension and axis. 
  Thanks (and thanks, Keith),
     Ken

-------------- next part --------------
An HTML attachment was scrubbed...
URL: http://mail.scipy.org/pipermail/numpy-discussion/attachments/20100406/345d9f86/attachment-0001.html