如何向量化实现DataFrame中的折叠运算以提升大数据集效率？

针对你描述的递推计算需求（D[0] = C[0]，D[n] = D[n-1] * A[n] + B[n]），以下是几种高性能的实现方案，优先采用向量化逻辑提升大数据集运算效率：

方法1：数学推导转化为向量化运算 #

通过展开递推公式，可将循环逻辑转化为累积乘积与加权求和的组合，完全利用Pandas/Numpy的向量化函数实现：

import pandas as pd
import numpy as np

# 构造示例DataFrame
df = pd.DataFrame({'A': [np.nan, 0.5, 0.5, 0.5, 0.5], 
                   'B': [np.nan, 3, 4, 1, 2], 
                   'C': [10, np.nan, np.nan, np.nan, np.nan]})

# 提取有效计算项：A从第1行开始，补1（不影响初始乘积）；B从第1行开始
A_vals = df['A'].fillna(1).iloc[1:].values
B_vals = df['B'].fillna(0).iloc[1:].values
initial_D = df['C'].iloc[0]

# 计算从后往前的累积乘积（对应每个B项后续的A乘积）
cum_prod = np.cumprod(A_vals[::-1])[::-1]
# 计算B项的加权和
weighted_sum = np.cumsum(B_vals * np.append(cum_prod[1:], 1))
# 拼接得到完整D列
D_col = np.concatenate([[initial_D], initial_D * cum_prod[0] + weighted_sum])

df['D'] = D_col

运行结果与示例完全一致：

A    B     C     D
0  NaN  NaN  10.0  10.0
1  0.5  3.0   NaN   8.0
2  0.5  4.0   NaN   8.0
3  0.5  1.0   NaN   5.0
4  0.5  2.0   NaN   4.5

方法2：Numba编译加速循环 #

如果递推逻辑复杂无法推导数学公式，用Numba将循环编译为机器码，性能接近纯向量化：

from numba import jit

@jit(nopython=True)
def compute_D(A, B, C):
    n = len(A)
    D = np.empty(n, dtype=np.float64)
    D[0] = C[0]
    for i in range(1, n):
        D[i] = D[i-1] * A[i] + B[i]
    return D

# 传入数组计算
df['D'] = compute_D(df['A'].values, df['B'].values, df['C'].values)

这种方法比纯Python循环快数十倍，适合复杂递推场景。

方法3：Pandas expanding.apply（简洁型） #

虽然本质仍为循环，但结合expanding API实现更简洁，性能优于纯Python循环：

def calc_d(window):
    if len(window) == 1:
        return window['C'].iloc[0]
    prev_d = window['D'].iloc[-2]
    curr_a = window['A'].iloc[-1]
    curr_b = window['B'].iloc[-1]
    return prev_d * curr_a + curr_b

df['D'] = df.expanding().apply(calc_d, axis=0)['D']

适合小数据集或快速原型开发，性能略逊于前两种方法。

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