2010/4/5 Ken Basye <kbas...@jhu.edu>:
> 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.
Maybe:
def select(A, B, axis):
# Extract incomplete index tuples:
argmax = a.argmax(axis = axis)
# Create the selection list to be handed over to B.__getitem__() ...
advanced_index = []
# Decompose the index tuples.
for dimi in xrange(argmax.shape[1]):
advanced_index.append(argmax[:, dimi])
# Insert the missing dimension.
advanced_index.insert(axis, numpy.arange(0, B.shape[axis]))
# Perform advanced (integer) selection ...
return B[advanced_index]
>>> a
array([[[0, 1],
[2, 3]],
[[4, 5],
[6, 7]]])
>>> select(a, a, 0)
array([3, 7])
>>> select(a, a, 1)
array([5, 7])
>>> select(a, a, 2)
array([6, 7])
It seems to work.
Friedrich
_______________________________________________
NumPy-Discussion mailing list
NumPy-Discussion@scipy.org
http://mail.scipy.org/mailman/listinfo/numpy-discussion
