How do I find the index for the minimum, second to minimum, etc. in a multi-dimensional list in Python?

Question

I currently have the following code that creates an nxnxn list of values. I am using the following code to find the index of the minimum value of the whole list, but I am also interested in finding the second to minimum, third to minimum and so on. Here is the code I am using right now: You can ignore the for loops and such as they all seem to be working properly.

    ind = np.unravel_index(C.argmin(), C.shape)

C is the list that contains the data. Is there a simple way to modify this line of code to find what I am looking for?

Answer 1

You could use np.argsort instead of argmin :

>>> import numpy as np
>>> arr = np.random.random((20, 20, 20))
>>> nsmallest = 3
>>> np.unravel_index(np.argsort(arr, axis=None)[:nsmallest], arr.shape)
(array([ 3, 16, 12], dtype=int64),
 array([ 2, 15, 15], dtype=int64),
 array([16, 11, 19], dtype=int64))

This is however transposed so the first column contains the index for the smallest item: (3, 2, 16) , the second the index for the second-to-smallest: (16, 15, 11) , ...

Or simply use a loop over the argsort result:

>>> argsorted = np.argsort(arr, axis=None)
>>> for i in range(nsmallest):
...     print(np.unravel_index(argsorted[i], arr.shape))
(3, 2, 16)
(16, 15, 11)
(12, 15, 19)