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

[Vincent Schut]

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! :-)