Pandas：不使用循环基于计算出的LAG值计算Rolling列

输入数据集 #

需求说明 #

按Group分组后，对Height≠0的行，用上一行的Rolling值减去当前行的Value填充Rolling列，要求尽可能避免使用循环。

输出数据集 #

实现方案 #

1. SQL 实现 #

利用窗口函数计算分组内的累积值，结合初始Rolling值反向推导后续行结果，无需循环：

WITH base_data AS (
    SELECT 
        "Group",
        Height,
        Value,
        Rolling,
        -- 计算分组内从初始行之后开始的Value累积和
        SUM(CASE WHEN Height = 0 THEN 0 ELSE Value END) OVER (PARTITION BY "Group" ORDER BY Height) AS cumulative_value
    FROM your_table
)
SELECT 
    "Group",
    Height,
    Value,
    -- 初始行保留原Rolling值，后续行用初始值减去累积和
    CASE 
        WHEN Height = 0 THEN Rolling 
        ELSE (SELECT Rolling FROM base_data b2 WHERE b2."Group" = b1."Group" AND b2.Height = 0) - cumulative_value 
    END AS Rolling
FROM base_data b1
ORDER BY "Group", Height;

2. Python Pandas 实现 #

通过分组累积求和与向量运算实现，避免循环：

import pandas as pd

# 构造输入数据
df = pd.DataFrame({
    'Group': ['A', 'A', 'A', 'B', 'B', 'B'],
    'Height': [0, 10, 20, 0, 10, 20],
    'Value': [2, 5, 3, 1, 4, 2],
    'Rolling': [100, None, None, 98, None, None]
})

# 获取每个分组的初始Rolling值
initial_rolling = df.loc[df['Height'] == 0].set_index('Group')['Rolling']
# 计算分组内的Value累积和（从初始行之后开始，所以shift(1)）
df['cum_sum'] = df.groupby('Group')['Value'].shift(1).cumsum().fillna(0)
# 填充Rolling列
df['Rolling'] = df.apply(
    lambda row: initial_rolling[row['Group']] - row['cum_sum'] if pd.isna(row['Rolling']) else row['Rolling'],
    axis=1
)
df.drop('cum_sum', axis=1, inplace=True)

print(df)

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