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.])
I tried several approaches with take() but now that I understand how
take() works when you give it an axis argument it seems like this isn't
going to do what I want. Still, it seems like there should be some
shortcut...
TIA,
Ken
_______________________________________________
NumPy-Discussion mailing list
NumPy-Discussion@scipy.org
http://mail.scipy.org/mailman/listinfo/numpy-discussion
