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

Question

例如，

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 代码对于只执行两个简单的算术并没有那么慢，而自定义函数方法是最快的。

Answer 1

你可以使用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。