python - 下面的代码有更快的实现吗？

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 循环，如果 a 非常大，它会相当慢。在 numpy 本身中有实现这个的技巧吗？

Answer 1

您可以将其表述为 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. 请参阅下面的评论，了解有关加快速度的更多想法。