如何高效运行该循环？分组id滚动计算的优化需求

问题背景 #

有如下格式的分组数据：

id    rr
    2.0   4
    2.0   5 
    2.0   4
    2.0   3
    5.0   4
    5.0   2 
    5.0   4
    5.0   1
    7.0   1
    7.0   2 
    7.0   3
    7.0   2

需要为每个id计算两个滚动统计参数：

• sum_rr2：窗口大小为3的滚动和减去当前行的rr值
• sum_rr3：窗口大小为4的滚动和减去当前行的rr值

原代码通过循环每个唯一id、切片数据再拼接结果的方式实现，但面对436个id、每个id对应8760行数据的场景，效率极低，原代码如下：

import pandas as pd 
sum_rr = pd.DataFrame(columns = {'id' ,'sum_rr2', 'sum_rr3'})

for i in id_cat.unique():
    test = data[data.id == i]
    test['sum_rr2'] = test.rr.rolling(3 , center = False).sum() - test.rr 
    test['sum_rr3'] = test.rr.rolling(4 , center = False).sum() - test.rr 
    sum_rr = sum_rr.append(test[['sum_rr3','sum_rr4','id']] , ignore_index = True)

sum_rr = sum_rr.reset_index(drop = True)

注：原代码中append时引用的sum_rr4应为笔误，实际应为sum_rr2

优化方案 #

使用pandas的groupby结合矢量化滚动操作替代循环，避免循环切片和重复append的性能损耗（append每次都会生成新DataFrame，内存和时间开销极大）。优化后的代码如下：

import pandas as pd

# 按id分组，对rr列计算滚动统计量
data['sum_rr2'] = data.groupby('id')['rr'].rolling(3, center=False).sum().reset_index(level=0, drop=True) - data['rr']
data['sum_rr3'] = data.groupby('id')['rr'].rolling(4, center=False).sum().reset_index(level=0, drop=True) - data['rr']

# 提取需要的列（按原代码需求调整列顺序）
sum_rr = data[['id', 'sum_rr2', 'sum_rr3']].copy()

优化点说明 #

• 分组矢量化操作：groupby('id')['rr'].rolling(...)直接对每个分组执行滚动计算，完全利用pandas的内部优化，比循环每个id的效率高一个数量级。
• 避免重复创建DataFrame：直接在原数据上计算新列，最后一次性提取结果，规避了append带来的频繁内存分配和拷贝。
• 修正笔误：将原代码中错误的sum_rr4修正为sum_rr2，保证结果列正确。

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