如何在Pandas中按ID分组，基于rolling列实现指定窗口的行滚动求和？

现有已按time字段排序的Pandas DataFrame，需按ID列分组（同一ID的rolling列值固定，不同ID可共用同一rolling值），为每行计算其**最后n行（含当前行）**的value列求和结果（窗口不足n行时，前面补0），最终得到expected_column。

示例数据 #

import pandas as pd

dct_data = {'ID': {0: 'a',
                   1: 'a',
                   2: 'a',
                   3: 'a',
                   4: 'a',
                   5: 'b',
                   6: 'b',
                   7: 'b',
                   8: 'b',
                   9: 'b'},
            'time': {0: '2022-12-23 14:56:00',
                     1: '2022-12-23 14:57:00',
                     2: '2022-12-23 14:58:00',
                     3: '2022-12-23 14:59:00',
                     4: '2022-12-23 15:00:00',
                     5: '2022-12-23 14:56:00',
                     6: '2022-12-23 14:57:00',
                     7: '2022-12-23 14:58:00',
                     8: '2022-12-23 14:59:00',
                     9: '2022-12-23 15:00:00'},
            'rolling': {0: 3, 1: 3, 2: 3, 3: 3, 4: 3, 5: 2, 6: 2, 7: 2, 8: 2, 9: 2},
            'value': {0: 19, 1: 14, 2: 14, 3: 32, 4: 16, 5: 0, 6: 6, 7: 1, 8: 4, 9: 3} }

df_test = pd.DataFrame(dct_data)

预期输出 #

ID  time                    rolling   value  expected_column
a   2022-12-23 14:56:00     3         19     19
a   2022-12-23 14:57:00     3         14     33
a   2022-12-23 14:58:00     3         14     47
a   2022-12-23 14:59:00     3         32     60
a   2022-12-23 15:00:00     3         16     62
b   2022-12-23 14:56:00     2         0      0
b   2022-12-23 14:57:00     2         6      6
b   2022-12-23 14:58:00     2         1      7
b   2022-12-23 14:59:00     2         4      5
b   2022-12-23 15:00:00     2         3      7

解决方法 #

利用groupby按ID分组，对每个分组先补全窗口所需的前置0，再计算固定窗口的滚动求和，最后匹配回原分组数据：

def calculate_rolling_sum(group):
    # 获取当前分组的窗口大小（同一ID的rolling值固定）
    window_size = group['rolling'].iloc[0]
    # 在value列前填充window_size-1个0，确保窗口不足时补0求和
    padded_values = pd.concat([pd.Series([0]*(window_size-1)), group['value']])
    # 计算滚动窗口求和
    rolling_sum = padded_values.rolling(window=window_size).sum()
    # 截取对应原分组长度的结果，赋值给新列
    group['expected_column'] = rolling_sum.iloc[window_size-1:].values
    return group

# 分组计算并更新DataFrame
df_test = df_test.groupby('ID', group_keys=False).apply(calculate_rolling_sum)

# 查看结果
print(df_test)

逻辑说明 #

• 同一ID的rolling值固定，直接取分组内第一个rolling值作为窗口大小
• 前置填充window_size-1个0，保证分组前几行也能凑够窗口数量求和（不足位置用0补充）
• 滚动求和后，截取从第window_size-1位开始的结果，正好对应原分组的每一行数据

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