Search code examples
pythonnumpymin

How to get indices of np.amin of 3d np array?


Given the following code:

import numpy as np

x = np.array([[[1, 2],
               [3, 4]],
              
              [[5, 6],
               [7, 8]], 
              
              [[3, 1],
               [1, 5]]])

x_min = np.amin(x, axis=0)
print(x_min)

The output (x_min) is

[[1 1]
 [1 4]]

Now I want to get the indices of dimension 0 of array x for the results of x_min array, it should be:

[[0 2]
 [2 0]]

Which function can I use to get this indices?


Solution

  • Try np.argmin: np.argmin(x, axis=0)