我想优化一个涉及递归公式的计算:
#a, b, c are arrays of shape (N, )
a[0] = c[0]
for i in range(1, N):
a[i] = a[i-1]*b[i] + c[i]
有没有一种方法可以通过优化的NumPy方法组合返回结果a
?比如cumprod
、cumsum
和dot
?
我想优化一个涉及递归公式的计算:
#a, b, c are arrays of shape (N, )
a[0] = c[0]
for i in range(1, N):
a[i] = a[i-1]*b[i] + c[i]
有没有一种方法可以通过优化的NumPy方法组合返回结果a
?比如cumprod
、cumsum
和dot
?
No, this sort of algorithms are hard to vectorize (using just numpy). I recommend to look at numba. Example:
import numba
import numpy as np
@numba.njit
def compute(a, b, c):
a[0] = c[0]
for i in range(1, len(a)):
a[i] = a[i - 1] * b[i] + c[i]
N = 10
a = np.empty(N, dtype="int64")
b = np.arange(1, N + 1, dtype="int64")
c = np.arange(10, N + 10, dtype="int64")
# note - when run for first time, numba compiles the function (this takes some time).
# Next runs numba uses already compiled function which is fast
compute(a, b, c)
print(a)
打印:
[ 10 31 105 433 2179 13089 91639 733129
6598179 65981809]