Question

我正在尝试优化一些代码，通过分析我注意到这个特定的循环需要花费很多时间。你能帮我写得更快吗？

import numpy as np

rows_a, rows_v, cols = (10, 15, 3)
a = np.arange(rows_a*cols).reshape(rows_a,cols)
v = np.arange(rows_v*cols).reshape(rows_v,cols)

c = 0
for i in range(len(v)):
    D = ((a-v[i])**2).sum(axis=-1)
    c += D.min()

print(c)

是否有任何numpy功能可以有效地做到这一点？

Answer 1

import numpy as np

rows_a, rows_v, cols = (10, 15, 3)
a = np.arange(rows_a*cols).reshape(rows_a,cols)
v = np.arange(rows_v*cols).reshape(rows_v,cols)

def using_loop():
    c = 0
    for i in range(len(v)):
        D = ((a-v[i])**2).sum(axis=-1)
        c += D.min()
    return c

def using_broadcasting():
    return ((a[:,np.newaxis,:]-v)**2).sum(axis=-1).min(axis=0).sum()

In [106]: %timeit using_loop()
1000 loops, best of 3: 233 µs per loop

In [107]: %timeit using_broadcasting()
10000 loops, best of 3: 29.1 µs per loop

In [108]: assert using_loop() == using_broadcasting()

使用NumPy时，通常有助于消除for-loops（如果可能）并使用在整个阵列上完成的操作来表达计算 - 或者至少在尽可能大的阵列上表达。通过这样做，您可以将更多工作卸载到使用C或Fortran编写的快速算法，而无需中间Python代码。

在原始代码中，D对于循环的每次迭代都具有形状(10,)。由于循环有15次迭代，如果我们可以将所有15次迭代中D的所有值一次表示为一个大数组，那么D将具有形状(10, 15)。事实上，我们可以这样做：

由于a的形状为(10,3)，a[:, np.newaxis, :]的形状为(10,1,3)。

使用NumPy broadcasting，因为v的形状为(15,3)，

a[:,np.newaxis,:]-v

具有形状(10,15,3)。平方，然后在最后一个轴上求和，得到一个形状(10, 15)的数组。这是新的D：

In [109]: ((a[:,np.newaxis,:]-v)**2).sum(axis=-1).shape
Out[109]: (10, 15)

一旦你有了D，剩下的计算就会自然而然。

两个numpy数组的每对列之间的差异（如何更有效地执行）？

1 个答案: