如何在 numpy 中对此计算进行矢量化

Question

如果 a 和 b 是大小合适的 numpy 数组，我如何矢量化以下计算？

total = a[0]
for ix in range(1, len(a)):
total = total*b[ix-1] + a[ix]

Answer 1

没关系，如果你做一些代数，有一个ufunc技巧是有效的。 在这种情况下，ufuncmultiply，诀窍accumulate.

c = np.r_[np.multiply.accumulate(b[0:-1][::-1])[::-1], 1]
total2 = np.sum(a * c)

这有什么作用：代数上，你求a[i]乘以b[i:]乘积for i in range(a.size)。 为此，翻转b并获取除最后一个数字之外的所有数字的运行乘积(假设a和b的长度相同(，然后将其翻转回来。 最后一个值应为 1，因为刚刚添加了最终的a值。

测试

a = np.random.randint(1, 10, 40)
b = 1 + np.random.rand(40)
total = a[0]
for ix in range(1, len(a)):
total = total*b[ix-1] + a[ix]
total
278443993.10494208
total2
278443993.10494208