[Numpy-discussion] Extracting values from one array corresponding to argmax elements in another array
Vincent Schut
schut@sarvision...
Tue Apr 6 05:51:12 CDT 2010
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! :-)
More information about the NumPy-Discussion
mailing list