Say I have a 3d matrix called A with shape (100,4,100). And also two vectors of 200 values each, created like this:

b = np.random.randint(0, 3, 200)
c = np.random.randint(0, 99, 200)

How can I efficiently find the argmax for each of A[b,c,:]?

I am able to do a for loop, which would look like this:

for i in range(200):
    np.argmax(A[b[i], c[i], :])

Do I have any more efficient options by avoiding the for loop?

CodePudding user response：

The solution as suggested by Divakar above is to use ndarray.argmax, so it would look like A[b,c].argmax(1)