有效地同时获得numpy.argmin和numpy.amin

时间:2016-05-19 16:08:05

标签: python performance numpy

是否可以通过一次调用numpy获得numpy.argmin和numpy.amin的结果?感谢。

2 个答案:

答案 0 :(得分:6)

您可以使用np.argmin获取与最小值对应的索引,然后使用NumPy's advanced indexing来获取最小值。

让我们使用2D数组来展示其用法。让A2D数组,假设我们有兴趣在axis=1找到最小索引和值。

因此,我们可以做 -

min_idx = np.argmin(A,axis=1)
min_val = A[np.arange(A.shape[0]),min_idx]

让我们为实例运行采用实际2D数组并验证结果 -

In [16]: A
Out[16]: 
array([[79, 97, 12, 54, 59],
       [44, 45, 42, 78, 32],
       [32, 41, 67, 60,  4],
       [24,  4, 85, 94, 65]])

In [17]: min_idx = np.argmin(A,axis=1)

In [18]: A[np.arange(A.shape[0]),min_idx] # Using min_idx & indexing
Out[18]: array([12, 32,  4,  4])

In [19]: np.amin(A,axis=1)                # Using np.amin to verify
Out[19]: array([12, 32,  4,  4])

运行时测试 -

In [26]: def original_app(A):
    ...:     min_idx = np.argmin(A,axis=1)
    ...:     min_val = np.amin(A,axis=1)
    ...:     return min_idx, min_val
    ...: 
    ...: def proposed_app(A):
    ...:     min_idx = np.argmin(A,axis=1)
    ...:     min_val = A[np.arange(A.shape[0]),min_idx]
    ...:     return min_idx, min_val
    ...: 

In [27]: A = np.random.randint(0,99,(4000,5000))

In [28]: %timeit original_app(A)
10 loops, best of 3: 70.9 ms per loop

In [29]: %timeit proposed_app(A)
10 loops, best of 3: 33.1 ms per loop

让我们更多地剖析时间 -

In [31]: min_idx = np.argmin(A,axis=1)

In [32]: %timeit np.argmin(A,axis=1)              # Used in both methods
10 loops, best of 3: 34.5 ms per loop

In [33]: %timeit np.amin(A,axis=1)                # Original approach
10 loops, best of 3: 37.3 ms per loop

In [34]: %timeit A[np.arange(A.shape[0]),min_idx] # Proposed approach
10000 loops, best of 3: 56 µs per loop

我们可以看到在最后一步使用高级索引获得了重大收益,并且花费的运行时间可以忽略不计。这允许几乎100%的运行时削减它!

答案 1 :(得分:0)

更新版本

一种非常干净和简单的方法是使用take_along_axis

参加:

In [1]: A = np.array([[[5, 1, 2, 3]], [[3, 5, 1, 7]]])

并申请:

In [2]: min_idx = np.argmin(A, axis=2)
In [3]: min_val = np.take_along_axis(A, min_idx[:,:,None], axis=2)[:,:,0]

就是这样! min_val等于np.amin(A, axis=2)


旧版本

如果您将解决方案从Divakar应用于超过2个维度,则在任何情况下所提供的解决方案都将无法使用。在解决了这个问题很长时间之后,我找到了一个似乎可行的解决方案(我在某些情况下对其进行了测试)。如果您发现此解决方案有任何问题,欢迎提出任何意见。

问题

那么首先是问题(与我遇到的问题类似)。尝试例如:

In [1]: A = np.array([[[5, 1, 2, 3]], [[3, 5, 1, 7]]])

在这种情况下,第二维仅为1(A的形状为(2,1,4))。然后

In [2]: np.amin(A, axis=2)

将导致

Out[2]: array([[1],
               [1]])

这是预期的。

另一方面,如果您申请

In [3]: min_idx = np.argmin(A, axis=2)
In [4]: A[np.arange(A.shape[0]), np.arange(A.shape[1]), min_idx]

你得到

Out[4]: array([[1, 5],
               [2, 1]])

这不等于amin中的In [2]解决方案。

旧解决方案

可能的是,在将min_idx(从In [3]中挤压出来之前,先将其用作索引,然后再对其进行重塑:

In [5]: sq_idx = np.squeeze(min_idx)
In [6]: min_pre = A[np.arange(A.shape[0]), np.arange(A.shape[1]), sq_idx]
In [7]: np.reshape(min_pre, (2,1))

最终导致:

Out[7]: array([[1],
               [1]])

我建议再采取两个步骤来优化解决方案。

(1)

为了简化数组的索引编制:

shp = [np.arange(i) for i in A.shape]

In [6]简化为:

min_pre = A[shp[0], shp[1], sq_idx]

(2)

为了概括重塑效果,In [7]可以替换为

np.reshape(min_pre, tuple(np.asarray(A.shape)[:-1]))

我希望它可以帮助某人:)