CODEWITHSUNDEEP
pythonnumpyargmax
Vitaly Olegovitch
Vitaly Olegovitch

Reputation: 3547

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

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?

Upvotes: 1

Views: 507

Answers (1)

I'mahdi
I'mahdi

Reputation: 24049

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])

Upvotes: 1

Related Questions