如何实现Rolling grouped cumulative sum？寻求非迭代高效方案

嘿，我来帮你搞定分组滚动累计求和的优化方案！迭代法虽然能跑，但数据量大的时候效率真的拉胯，用Pandas的内置函数能直接搞定，效率高多了。先给你梳理下具体实现：

先明确源数据与期望输出 #

假设你的源数据是这样的（用Pandas DataFrame示例）：

import pandas as pd

df = pd.DataFrame({
    'group': ['A', 'A', 'A', 'B', 'B', 'B', 'B'],
    'value': [10, 20, 30, 5, 15, 25, 35]
})

如果滚动窗口大小设为3，期望的输出应该是：

最优实现：Pandas内置函数一行搞定 #

直接利用Pandas的分组滚动计算能力，底层是C优化的，比手动迭代快几十倍甚至上百倍，代码还简洁：

# 设置滚动窗口大小，可根据需求调整
window_size = 3

# 核心代码：分组+滚动求和
df['rolling_cumsum'] = df.groupby('group')['value'].rolling(
    window=window_size, 
    min_periods=1  # 窗口不足指定大小时仍计算，比如每组前几个数据
).sum().reset_index(level=0, drop=True)

参数解释：

• groupby('group')：按指定分组列拆分数据
• rolling(window=window_size, min_periods=1)：创建滚动窗口，min_periods=1保证即使窗口没填满（比如每组前2条数据）也会计算累计和
• sum()：对窗口内的数值求和
• reset_index(level=0, drop=True)：移除分组产生的额外索引，让结果能和原DataFrame完美合并

为什么不推荐迭代法？ #

手动迭代的问题很明显：

• 效率极低：循环每个分组、每个行是纯Python层面的操作，数据量过万时速度会慢到难以接受；而Pandas的内置函数是底层C实现，计算效率碾压手动迭代
• 代码冗余易出错：手动处理分组边界、窗口大小、索引对齐很容易写出bug，维护成本高

纯Python替代方案（无Pandas场景） #

如果你的项目没法用Pandas，也可以用标准库的itertools来优化，比手动写循环高效：

from itertools import groupby, accumulate
import operator

# 先确保数据按分组排序（如果源数据未排序）
sorted_data = sorted(df.to_dict('records'), key=lambda x: x['group'])
window_size = 3

result = []
for group_key, group_items in groupby(sorted_data, key=lambda x: x['group']):
    values = [item['value'] for item in group_items]
    # 先计算前缀和，再推导滚动窗口和
    prefix_sums = list(accumulate(values, operator.add))
    rolling_sums = []
    for i in range(len(prefix_sums)):
        if i < window_size - 1:
            rolling_sums.append(prefix_sums[i])
        else:
            rolling_sums.append(prefix_sums[i] - (prefix_sums[i - window_size] if i >= window_size else 0))
    # 将结果映射回原数据
    group_records = [item for item in sorted_data if item['group'] == group_key]
    for rec, rs in zip(group_records, rolling_sums):
        rec['rolling_cumsum'] = rs
        result.append(rec)

# 转换回你需要的格式（比如列表或DataFrame）
final_df = pd.DataFrame(result)

不过这个方案还是不如Pandas高效，所以优先推荐用Pandas的实现。

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