如何向量化实现Pandas中递推式迭代计算函数？

你的递推计算可以通过数学展开+累积运算组合实现完全向量化，核心是先把循环递推式转换成可通过cumprod和cumsum计算的形式。

递推式数学展开 #

原函数的递推逻辑是：

• $r_0 = foo[0]$
• $r_i = bar[i] + baz[i] \times r_{i-1} \quad (i \geq 1)$

将递推式展开后可得：
$$
r_i = foo[0] \times \prod_{k=1}^i baz[k] + \sum_{k=1}^i \left( bar[k] \times \prod_{m=k+1}^i baz[m] \right)
$$

进一步变形（提取公共乘积项）：
$$
r_i = \prod_{k=1}^i baz[k] \times \left( foo[0] + \sum_{k=1}^i \frac{bar[k]}{\prod_{m=1}^k baz[m]} \right)
$$

这个变形后的式子就可以用cumprod（累积乘积）和cumsum（累积求和）来实现向量化计算。

正确的向量化实现 #

import pandas as pd

def calculate_result_vectorized(foo: pd.Series, bar: pd.Series, baz: pd.Series) -> pd.Series:
    # 构造累积乘积数组：prod[0]=1，prod[i] = baz[1] * baz[2] * ... * baz[i]
    prod = pd.Series(index=foo.index, dtype=float)
    prod.iloc[0] = 1.0
    if len(foo) > 1:
        prod.iloc[1:] = baz.iloc[1:].cumprod().values
    
    # 构造求和项的每一项，首项设为0（因为求和从i=1开始）
    term = bar / prod
    term.iloc[0] = 0.0
    
    # 计算累积和
    cum_term = term.cumsum()
    
    # 计算最终结果
    result = prod * (foo.iloc[0] + cum_term)
    # 确保首项与原函数一致（计算逻辑已保证，此处可省略）
    result.iloc[0] = foo.iloc[0]
    
    return result

验证结果 #

用你的示例数据测试：

foo = pd.Series([1, 2, 3, 4])
bar = pd.Series([2, 3, 4, 5])
baz = pd.Series([3, 4, 5, 6])

print(calculate_result_vectorized(foo, bar, baz))

输出：

0      1.0
1      7.0
2     39.0
3    239.0
dtype: float64

与原迭代函数的结果完全一致。

性能优势 #

向量化实现避免了Python级别的循环，利用Pandas/NumPy的底层C优化运算，在数据量大时（比如百万级数据）性能会比迭代函数提升数十倍甚至上百倍，非常适合频繁调用的场景。

内容的提问来源于stack exchange，提问作者Asclepius