逐行访问前置值的循环向量化优化方案咨询

AI 小北

2026-6-19

Pandas递推计算的向量化优化方案（替代iterrows循环）

你的需求是基于前一行的column1和column2值递推更新当前行的column2，原生iterrows循环在数据量较大时效率极低，以下是几种高效的优化方案：

方案一：数学推导+纯向量化计算

先拆解递推逻辑：

初始值：column2[0] = 0
递推公式：column2[n] = DECAY * (column1[n-1] + column2[n-1])

展开后可发现，column2[n]本质是前n个column1元素分别乘以DECAY的幂次后的累加，公式为：
column2[k] = DECAY*column1[k-1] + DECAY²*column1[k-2] + ... + DECAYᵏ*column1[0]

基于这个规律，用numpy的向量化运算可以直接计算：

import pandas as pd
import numpy as np

DECAY = 0.9
df = pd.DataFrame({'column1': [1, 2, 3, 4, 5]})

n_rows = len(df)
# 生成权重序列：DECAY^1, DECAY^2, ..., DECAY^(n-1)
weights = DECAY ** np.arange(1, n_rows)
# 用卷积计算加权和，对应每行column2的值
df['column2'] = 0.0
df.loc[1:, 'column2'] = np.convolve(df['column1'].values[:-1], weights, mode='valid')

# 输出结果
print(df)

方案二：Numba JIT加速循环

如果递推逻辑复杂、无法推导通项，用Numba对循环进行JIT编译，能把原生Python循环的速度提升数百倍：

import pandas as pd
import numba as nb

DECAY = 0.9
df = pd.DataFrame({'column1': [1, 2, 3, 4, 5]})

@nb.jit(nopython=True)
def compute_column2(c1_array):
    n = len(c1_array)
    c2_array = np.zeros(n, dtype=np.float64)
    for i in range(1, n):
        c2_array[i] = (c1_array[i-1] + c2_array[i-1]) * DECAY
    return c2_array

df['column2'] = compute_column2(df['column1'].values)
print(df)