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

Question

我有一个一维的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？

Answer 1

您可以将此短语称为外部产品:

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.

正如您所看到的,这可以提供~1000倍的速度n=100000.

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

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

PS表达这个想法的另一种方法是使用元素乘法与广播(基于@Scott Gigante的建议):

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

PPS请参阅以下评论,了解有关加快这一进程的进一步想法.