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?
b[row , column], with
argmax(0) you find columns, for adding row you can use
>>> b[range(len(a)) , a.argmax(0)] array([6, 8])