I have an array a with shape (18,4096,4096).

And I want to do like these:

max_value = np.max(a,0)
index = np.argmax(a,0)

max_value and index are both array with shape (4096, 4096), and I think calling both np.max and np.argmax has some useless cost.

And I know if a is a 1D array, I can do like this:

index = np.argmax(a,0)
max_value = a[index]

But I can't do like this when a is a 3D array. Is there any efficient way doing this?

CodePudding user response：

This should work:

i, j = np.meshgrid(np.arange(a.shape[1]), np.arange(a.shape[2]), indexing='ij')
max_value = a[index[i,j],i,j]

or following @hpaulj's suggestion (~2x faster):

max_value  = np.take_along_axis(a, index[None], 0)

CodePudding user response：

This works for me:

index = np.unravel_index(np.argmax(a, axis=None), a.shape)
max_value=a[index]