numpy:获取二维数组最小值的列和行索引

发布于 2021-01-29 16:15:48

例如,

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代码执行两个简单的算术并没有那么慢,而自定义函数方法却可以做到最快。

关注者
0
被浏览
160
1 个回答
  • 面试哥
    面试哥 2021-01-29
    为面试而生,有面试问题,就找面试哥。

    您可以使用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的原始方法的一种方法是使用divmod

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

    对它们进行计时,您会发现额外的速度点虽然不多,但仍然有所改进。

    %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



知识点
面圈网VIP题库

面圈网VIP题库全新上线,海量真题题库资源。 90大类考试,超10万份考试真题开放下载啦

去下载看看