Python高效实现行末值作为下行起始值的滚动计算

需要在超大规模数据上执行滚动计算，避免嵌套for循环的性能问题，仅使用标准Python（含Pandas）。具体需求：

• 对pd.DataFrame按key1和key2分组（如AX、AY、BX组）
• 核心计算逻辑：end = start + plus - minus，且每行的end作为下一行的start
• 数值可为负、零或正，需避免内存溢出

原始数据df1及预期结果df2如下：

import pandas as pd 
import numpy as np

df1 = pd.DataFrame(np.array([["A", "X", 3,6,4,0], ["A", "X", 0,2,10,0], ["A", "X", 0,9,3,0], ["A", "Y", 8,3,1,0], ["A", "Y", 0,2,3,0], ["B", "X", 4,4,2,0], ["B", "X", 0,1,0,0]]),
                   columns=['key1', 'key2', 'start', 'plus', 'minus', 'end'])

print("原始数据df1:")
print(df1)
#   key1 key2 start plus minus end
# 0    A    X     3    6     4   0
# 1    A    X     0    2    10   0
# 2    A    X     0    9     3   0
# 3    A    Y     8    3     1   0
# 4    A    Y     0    2     3   0
# 5    B    X     4    4     2   0
# 6    B    X     0    1     0   0

df2 = pd.DataFrame(np.array([["A", "X", 3,6,4,5], ["A", "X", 5,2,10,-3], ["A", "X", -3,9,3,3], ["A", "Y", 8,3,1,10], ["A", "Y", 10,2,3,9], ["B", "X", 4,4,2,2], ["B", "X", 2,1,0,3]]),
                   columns=['key1', 'key2', 'start', 'plus', 'minus', 'end'])

print("\n预期结果df2:")
print(df2)
#   key1 key2 start plus minus end
# 0    A    X     3    6     4   5
# 1    A    X     5    2    10  -3
# 2    A    X    -3    9     3   3
# 3    A    Y     8    3     1  10
# 4    A    Y    10    2     3   9
# 5    B    X     4    4     2   2
# 6    B    X     2    1     0   3

步骤1：转换数值列类型 #

原始数据的数值列存储为字符串，先转为数值型以支持计算：

df = df1.copy()
for col in ['start', 'plus', 'minus', 'end']:
    df[col] = pd.to_numeric(df[col])

步骤2：计算分组内的净增量 #

先推导每行的净增量delta = plus - minus，将滚动依赖逻辑转化为累积和计算：

df['delta'] = df['plus'] - df['minus']

步骤3：分组计算累积增量与目标列 #

利用Pandas内置的分组累积运算替代循环，实现高效计算：

# 计算分组内从第一行到当前行的delta累积和
df['cum_delta'] = df.groupby(['key1', 'key2'])['delta'].cumsum()

# 获取每个分组的初始start值
df['initial_start'] = df.groupby(['key1', 'key2'])['start'].transform('first')

# 计算end列：初始start + 累积delta
df['end'] = df['initial_start'] + df['cum_delta']

# 更新start列：分组第一行保留原start，其余行取上一行的end
df['start'] = df.groupby(['key1', 'key2'])['end'].shift(1).fillna(df['initial_start'])

# 清理临时辅助列
df = df.drop(columns=['delta', 'cum_delta', 'initial_start'])

验证结果 #

将处理后的数据与预期结果对比：

# 先转换df2的数值列类型
df2_numeric = df2.copy()
for col in ['start', 'plus', 'minus', 'end']:
    df2_numeric[col] = pd.to_numeric(df2_numeric[col])

print(df.equals(df2_numeric))  # 输出：True

• 性能高效：基于Pandas向量化运算和分组累积，避免循环遍历，处理超大规模数据时性能提升显著
• 内存友好：所有计算基于原DataFrame列操作，无需额外开辟大量临时空间，降低内存溢出风险
• 逻辑简洁：通过数学转换将滚动依赖转化为累积和计算，符合Pandas的向量化编程范式

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