使用 NumPy 优化 a[i] = a[i-1]*b[i] + c[i] 的迭代计算

Question

我想优化一个涉及递归公式的计算：

#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

？

Answer 1

不，这种算法很难矢量化（仅使用 numpy）。我建议看看numba。示例：

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")

compute(a, b, c)
print(a)

打印：

[      10       31      105      433     2179    13089    91639   733129
  6598179 65981809]