以下代码是否有更快的实现?

时间:2018-10-23 02:00:47

标签: python performance loops numpy vectorization

我有一个一维numpy数组,它的大小很大。对于数组的每个条目,我需要生成一个线性间隔的子数组,直到该条目值。这是我的例子。

import numpy as np
a = np.array([2, 3])
b = np.array([np.linspace(0, i, 4) for i in a])

在这种情况下,存在大小为4的线性空间。上面代码中的最后一条语句涉及for loop,如果a非常大,则该速度很慢。是否可以在numpy本身中实现此目的?

1 个答案:

答案 0 :(得分:9)

您可以将其表达为outer product

In [37]: a = np.arange(100000)

In [38]: %timeit np.array([np.linspace(0, i, 4) for i in a])
1 loop, best of 3: 1.3 s per loop

In [39]: %timeit np.outer(a, np.linspace(0, 1, 4))
1000 loops, best of 3: 1.44 ms per loop

想法是采用一个linspace单位,然后按a的每个元素分别缩放。

如您所见,n=100000的速度提高了约1000倍。

为完整起见,我将提到此代码的舍入属性与原始版本略有不同(在实际应用中可能不是问题):

In [52]: np.max(np.abs(np.array([np.linspace(0, i, 4) for i in a]) -
    ...:               np.outer(a, np.linspace(0, 1, 4))))
Out[52]: 1.4551915228366852e-11

P。 S.表达此想法的另一种方法是通过广播使用逐元素乘法(基于@Scott Gigante的建议):

In [55]: %timeit a[:, np.newaxis] * np.linspace(0, 1, 4)
1000 loops, best of 3: 1.48 ms per loop

P。 P. S.请参阅下面的评论,以获取进一步加快此速度的想法。