使用另一个数组在 Numpy ndarray 中选择某些索引

Question

I'm trying to mark the value and indices of max values in a 3D array, getting the max in the third axis. Now this would have been obvious in a lower dimension:

argmaxes=np.argmax(array)
maximums=array[argmaxes]

but NumPy doesn't understand the second syntax properly for higher than 1D. Let's say my 3D array has shape (8,8,250). argmaxes=np.argmax(array,axis=-1) would return a (8,8) array with numbers between 0 to 250. Now my expected output is an (8,8) array containing the maximum number in the 3rd dimension. I can achieve this with maxes=np.max(array,axis=-1) but that's repeating the same calculation twice (because I need both values and indices for later calculations) I can also just do a crude nested loop:

for i in range(8):
   for j in range(8):
      maxes[i,j]=array[i,j,argmaxes[i,j]]

But is there a nicer way to do this?

Answer 1

You can use advanced indexing. This is a simpler case when shape is (8,8,3) :

arr = np.random.randint(99, size=(8,8,3))
x, y = np.indices(arr.shape[:-1])
arr[x, y, np.argmax(array,axis=-1)]

Sample run:

>>> x
array([[0, 0, 0, 0, 0, 0, 0, 0],
       [1, 1, 1, 1, 1, 1, 1, 1],
       [2, 2, 2, 2, 2, 2, 2, 2],
       [3, 3, 3, 3, 3, 3, 3, 3],
       [4, 4, 4, 4, 4, 4, 4, 4],
       [5, 5, 5, 5, 5, 5, 5, 5],
       [6, 6, 6, 6, 6, 6, 6, 6],
       [7, 7, 7, 7, 7, 7, 7, 7]])
>>> y
array([[0, 1, 2, 3, 4, 5, 6, 7],
       [0, 1, 2, 3, 4, 5, 6, 7],
       [0, 1, 2, 3, 4, 5, 6, 7],
       [0, 1, 2, 3, 4, 5, 6, 7],
       [0, 1, 2, 3, 4, 5, 6, 7],
       [0, 1, 2, 3, 4, 5, 6, 7],
       [0, 1, 2, 3, 4, 5, 6, 7],
       [0, 1, 2, 3, 4, 5, 6, 7]])    
>>> np.argmax(arr,axis=-1)    
array([[2, 1, 1, 2, 0, 0, 0, 1],
       [2, 2, 2, 1, 0, 0, 1, 0],
       [1, 2, 0, 1, 1, 1, 2, 0],
       [1, 0, 0, 0, 2, 1, 1, 0],
       [2, 0, 1, 2, 2, 2, 1, 0],
       [2, 2, 0, 1, 1, 0, 2, 2],
       [1, 1, 0, 1, 1, 2, 1, 0],
       [2, 1, 1, 1, 0, 0, 2, 1]], dtype=int64)

This is a visual example of array to help to understand it better: