On Tue, Apr 6, 2010 at 9:22 AM, Ken Basye <kbas...@jhu.edu> wrote: > From: Vincent Schut <sc...@sarvision.nl> > > 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! :-) > > > 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
an alternative to Vincent's general solution, if you have unique max or want all argmax is using a mask >>> a array([[ 100., 0., 0.], [ 0., 100., 100.], [ 0., 0., 0.]]) >>> b array([[0, 1, 2], [3, 4, 5], [6, 7, 8]]) >>> ax=0; b[a==np.expand_dims(a.max(ax),ax)] array([0, 4, 5]) >>> ax=1; b[a==np.expand_dims(a.max(ax),ax)] array([0, 4, 5, 6, 7, 8]) >>> aa=np.eye(3) >>> ax=1; b[aa==np.expand_dims(aa.max(ax),ax)] array([0, 4, 8]) >>> ax=0; b[aa==np.expand_dims(aa.max(ax),ax)] array([0, 4, 8]) Josef > > > _______________________________________________ > NumPy-Discussion mailing list > NumPy-Discussion@scipy.org > http://mail.scipy.org/mailman/listinfo/numpy-discussion > > _______________________________________________ NumPy-Discussion mailing list NumPy-Discussion@scipy.org http://mail.scipy.org/mailman/listinfo/numpy-discussion