Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

Numpy: get the column and row index of the minimum value of a 2D array

For example,

x = array([[1,2,3],[3,2,5],[9,0,2]])
some_func(x) gives (2,1)

I know one can do it by a custom function:

def find_min_idx(x):
    k = x.argmin()
    ncol = x.shape[1]
    return k/ncol, k%ncol

However, I am wondering if there's a numpy built-in function that does this faster.

Thanks.

EDIT: thanks for the answers. I tested their speeds as follows:

%timeit np.unravel_index(x.argmin(), x.shape)
#100000 loops, best of 3: 4.67 µs per loop

%timeit np.where(x==x.min())
#100000 loops, best of 3: 12.7 µs per loop

%timeit find_min_idx(x) # this is using the custom function above
#100000 loops, best of 3: 2.44 µs per loop

Seems the custom function is actually faster than unravel_index() and where(). unravel_index() does similar things as the custom function plus the overhead of checking extra arguments. where() is capable of returning multiple indices but is significantly slower for my purpose. Perhaps pure python code is not that slow for doing just two simple arithmetic and the custom function approach is as fast as one can get.

like image 335
N. H. Avatar asked May 12 '15 01:05

N. H.


People also ask

How do you find the index of the minimum value in a NumPy array?

The numpy argmin() function takes arr, axis, and out as parameters and returns the array. To find the index of a minimum element from the array, use the np. argmin() function.


2 Answers

You can use np.unravel_index

print(np.unravel_index(x.argmin(), x.shape))
(2, 1)
like image 136
Padraic Cunningham Avatar answered Oct 15 '22 12:10

Padraic Cunningham


You may use np.where:

In [9]: np.where(x == np.min(x))
Out[9]: (array([2]), array([1]))

Also as @senderle mentioned in comment, to get values in an array, you can use np.argwhere:

In [21]: np.argwhere(x == np.min(x))
Out[21]: array([[2, 1]])

Updated:

As OP's times show, and much clearer that argmin is desired (no duplicated mins etc.), one way I think may slightly improve OP's original approach is to use divmod:

divmod(x.argmin(), x.shape[1])

Timed them and you will find that extra bits of speed, not much but still an improvement.

%timeit find_min_idx(x)
1000000 loops, best of 3: 1.1 µs per loop

%timeit divmod(x.argmin(), x.shape[1])
1000000 loops, best of 3: 1.04 µs per loop

If you are really concerned about performance, you may take a look at cython.

like image 20
Anzel Avatar answered Oct 15 '22 14:10

Anzel