如何用向量化方法实现Pandas DataFrame中v列的递推计算？

给定初始结构的pd.DataFrame：

n    v
    0    1  0.0
    1    2  0.0
    2    3  0.0
    3    4  0.0
    4    5  0.0
    5    6  0.0

需要计算v列的值，每个单元格的计算规则为：

当前v值 = 当前n值 + 前一个v值 × 固定系数（0.5）

计算示例如下：

## sample calculation table :~
n   v[i]      n  + v[i-1] * fixed factor 
1   1.0     = 1  + 0.0    * 0.5
2   2.5     = 2  + 1.0    * 0.5   
3   4.25    = 3  + 2.5    * 0.5
4   6.125   = 4  + 4.25   * 0.5
5   8.0625  = 5  + 6.125  * 0.5
6  10.03125 = 6  + 8.063  * 0.5

目前已通过逐行循环实现该逻辑，代码如下：

df = pd.DataFrame({'n':[1,2,3,4,5,6]})
df['v'] = 0.0

def fnv(nu, vu):
    return nu + vu * 0.5

for i in range(0, df.shape[0]):
    df.v.at[i] = fnv(df.n.at[i], df.v.at[i-1] if i>0 else 0.0)

运行结果符合预期，但循环效率较低，希望找到向量化的实现方法（避免逐行循环）。

方法1：利用递推公式的通项展开 #

观察递推规则v[i] = n[i] + 0.5*v[i-1]，可将其展开为加权求和形式：

• v[0] = n[0]
• v[1] = n[1] + 0.5*n[0]
• v[2] = n[2] + 0.5*n[1] + 0.5²*n[0]
• ...
• v[k] = sum_{i=0}^k n[i] * (0.5)^(k-i)

基于这个通项，用numpy广播机制实现向量化计算：

import pandas as pd
import numpy as np

df = pd.DataFrame({'n': [1,2,3,4,5,6]})

# 生成系数矩阵：第k行第i列的系数为0.5^(k-i)（仅i<=k时有意义）
coeffs = 0.5 ** np.arange(len(df))[::-1].reshape(-1, 1)
# 计算每一行的加权和得到v列
df['v'] = np.dot(coeffs, df['n'].values)

# 保留两位小数查看结果
print(df.round(2))

运行结果：

n      v
0  1   1.00
1  2   2.50
2  3   4.25
3  4   6.12
4  5   8.06
5  6  10.03

方法2：结合pandas expanding与向量化操作 #

用pandas的expanding获取每个位置的前序n值，再结合权重数组计算求和：

import pandas as pd
import numpy as np

df = pd.DataFrame({'n': [1,2,3,4,5,6]})
factor = 0.5  # 固定系数

# 生成对应位置的权重数组
weights = factor ** np.arange(len(df))[::-1]
# 对每个expanding窗口内的n值，乘以对应权重后求和
df['v'] = df['n'].expanding().apply(lambda x: (x * weights[-len(x):]).sum(), raw=True)

print(df.round(2))

该方法本质与方法1一致，更贴合pandas的API使用习惯。

方法3：用numba加速循环（准向量化） #

如果DataFrame行数极多，通项展开的内存占用过大，可使用numba对循环进行JIT编译，效率接近纯向量化：

import pandas as pd
from numba import jit

df = pd.DataFrame({'n': [1,2,3,4,5,6]})
df['v'] = 0.0

@jit(nopython=True)
def compute_v(n_arr, v_arr, factor):
    v_prev = 0.0
    for i in range(len(n_arr)):
        v_arr[i] = n_arr[i] + v_prev * factor
        v_prev = v_arr[i]

compute_v(df['n'].values, df['v'].values, 0.5)
print(df.round(2))

numba会将循环编译为机器码，效率远超原生Python循环。

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