Getting argmax indexes from one ndarray and use them to get values from another array-CodePudding

I am taking the argmax from one ndarray and I have to use it to get elements from another ndarray.

An example usage of this is "get the weights of the tallest people". Where the heights are in one ndarray and the weights are in another.

The code that I have now is:

a = np.array([[10, 20], [15, 5]])
print(a)
b = np.array([[7, 6], [8, 9]])
a_argmax = a.argmax(axis=0)
print(a_argmax)
print(b)
for i, m in enumerate(a_argmax):
    print(b[i, m])

Is it possible to achieve the same result in one shot, without a loop?

CodePudding user response：

You need b[row , column], with argmax(0) you find columns, for adding row you can use range(len(a)).

>>> b[range(len(a)) , a.argmax(0)]
array([6, 8])