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]) _______________________________________________ NumPy-Discussion mailing list NumPy-Discussion@scipy.org http://mail.scipy.org/mailman/listinfo/numpy-discussion
