Numpy:获取2D数组最小值的列和行索引

时间:2015-05-12 01:26:19

标签: python arrays numpy

例如,

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

我知道可以通过自定义函数来实现:

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

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

感谢。

编辑:谢谢你的回答。我测试了它们的速度如下:

%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代码并不是那么简单,只做两个简单的算术,自定义函数方法就像人们可以获得的那样快。

2 个答案:

答案 0 :(得分:11)

您可以使用np.where

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

同样@senderle在评论中提到,要获取数组中的值,您可以使用np.argwhere

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

更新

正如OP的时代所表明的那样,并且更加清楚argmin是理想的(没有重复的分钟等),我认为可能会略微改进OP的原始方法的一种方法是使用{{ 1}}:

divmod

定时他们,你会发现额外的速度,不多但仍然是一种改进。

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

如果您真的关心性能,可以看一下 cython

答案 1 :(得分:8)

您可以使用np.unravel_index

print(np.unravel_index(x.argmin(), x.shape))
(2, 1)