如何在 Pandas 中基于列指定窗口大小实现高效滚动求和

本文介绍一种无需显式循环、利用 pd.factorize 与 numpy.vstack 高效实现「每行使用不同窗口长度」的滚动求和方法，适用于 column b 动态定义 column a 滚动窗口大小的场景。

在 Pandas 里做滚动计算，最常见的需求是固定窗口大小，比如 rolling(3) 直接往前推三行就行。但现实业务总爱出一些“变数”——比方说，每一行的回溯周期得由另一列来指定：第 i 行要算 column A 从第 (i - window + 1) 行到第 i 行的和，而 window 却等于 column B 里当前行的值。这种“变长窗口”用 .rolling() 是没法一步到位的，但如果老老实实用 for 循环逐行算，数据一多就容易卡。今天要说的是一个向量化的小技巧，能绕过循环，跑起来快得多。

核心思路其实很直接：先把 column B 里所有不同的窗口长度找出来，对每个长度单独算一次完整的滚动求和，最后按原行对应的窗口类型“挑”出正确的结果就行。具体拆成四步来看：

1. 用 pd.factorize(df['B']) 把 column B 转成整数编码 idx 和唯一值数组 vals；
2. 针对每个唯一窗口长度 v，算一遍 df['A'].rolling(v, min_periods=1).sum()，得到长度跟原 DataFrame 一样的 Series；
3. 把这些 Series 纵向堆成一个二维数组，形状是 (len(vals), len(df))；
4. 用 idx 做行索引，配合 np.arange(len(df)) 做列索引，高级索引一出，每行的结果就精准到位了。

import pandas as pd
import numpy as np

# 构造示例数据
df = pd.DataFrame({
    'A': [1, 2, 1, 3, 2],
    'B': [1, 2, 3, 2, 4]
})

# 高效变长滚动求和
idx, vals = pd.factorize(df['B'])
df['C'] = np.vstack([
    df['A'].rolling(v, min_periods=1).sum()
    for v in vals
])[idx, np.arange(len(df))]

print(df)

运行结果是这样的：

   A  B    C
0  1  1  1.0
1  2  2  3.0
2  1  3  4.0
3  3  2  4.0
4  2  4  8.0

这里有几个要点值得留意：

• 参数 min_periods=1 保证了当窗口超出索引范围时（比如前几行窗口很大），依然能返回有效值，而不是直接给 NaN。如果业务要求严格满足窗口长度，可以改成 min_periods=v，然后自己处理 NaN；
• 这个方法的时间复杂度是 O(k × n)，k 是 column B 里不同窗口的个数，n 是总行数。只要 k 不算太大（比如几十种常见窗口），效率要远优于逐行 .apply(lambda x: ...)；
• 如果还想同时算均值、标准差，可以改用 .agg(['sum', 'mean']) 配合 np.dstack 来扩展，一次滚动对象的计算就搞定多项聚合。

这套方案在 Pandas 生态里算是一个处理“动态窗口滚动计算”的推荐实践，既保留了代码的可读性，性能上也不含糊。