On 04/05/2010 06:06 PM, Keith Goodman wrote: > On Mon, Apr 5, 2010 at 8:44 AM, Ken Basye<kbas...@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! :-) _______________________________________________ NumPy-Discussion mailing list NumPy-Discussion@scipy.org http://mail.scipy.org/mailman/listinfo/numpy-discussion
