numpy:获取二维数组最小值的列和行索引 [英] Numpy: get the column and row index of the minimum value of a 2D array

问题描述

例如，

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

但是，我想知道是否有一个numpy内置函数可以更快地完成此任务.

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

谢谢.

感谢您的回答.我测试了他们的速度，如下所示:

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

似乎自定义函数实际上比unravel_index()和where()更快. unravel_index()的功能与自定义函数相似，另外还需要检查额外参数的开销. where()能够返回多个索引，但是对于我而言，它的速度明显慢一些.纯粹的python代码执行两个简单的算术并没有那么慢，而自定义函数的方法就可以做到最快.

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.

推荐答案

您可以使用np.where:

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

也如注释中提到的@senderle一样，要获取数组中的值，您可以使用np.argwhere:

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]])

已更新:

正如OP的时间所显示的，并且更清楚地希望argmin是可取的(没有重复的分钟数等)，我认为可以稍微改善OP的原始方法的一种方法是使用divmod:

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

如果您真的很担心性能，可以看看 cython .

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

这篇关于numpy:获取二维数组最小值的列和行索引的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！