本文介绍一种无需显式循环、利用 pd.factorize 与 numpy.vstack 高效实现「每行使用不同窗口长度」的滚动求和方法，适用于 column b 动态定义 column a 滚动窗口大小的场景。
本文介绍一种无需显式循环、利用 pd.factorize 与 numpy.vstack 高效实现「每行使用不同窗口长度」的滚动求和方法，适用于 column b 动态定义 column a 滚动窗口大小的场景。
在 Pandas 中，标准的 rolling(window=n) 要求窗口大小为固定整数，而实际业务中常需根据另一列（如 Column B）动态指定每行的回溯周期——例如第 i 行的滚动和应取 df.loc[i-window+1:i, 'A'] 的和，其中 window = df.loc[i, 'B']。这种「变长窗口」无法直接用 .rolling() 实现，但可通过向量化技巧规避 Python 循环，显著提升性能。
核心思路是：对 Column B 中出现的所有唯一窗口长度分别计算完整滚动结果，再按原行顺序索引选取对应值。具体步骤如下：
- 使用 pd.factorize(df['B']) 将 Column B 映射为整数编码 idx 和唯一值数组 vals；
- 对每个唯一窗口长度 v，计算 df['A'].rolling(v, min_periods=1).sum()，得到形状为 (len(df), ) 的 Series；
- 将所有结果堆叠为二维数组 np.vstack([...])，其形状为 (len(vals), len(df))；
- 利用 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 确保首行即使窗口超出范围也能返回有效值（如窗口为 3 但仅前 1 行存在时，返回该行值）；若需严格满足窗口长度，可设 min_periods=v 并处理 NaN；
- 该方法时间复杂度为 O(k × n)，k 为 Column B 唯一值个数，n 为总行数——当 k 较小时（如几十种常见窗口）效率远优于逐行 apply(lambda x: ...)；
- 若需同时计算多种聚合（如均值、标准差），可用 agg(['sum', 'mean']) + np.dstack 扩展，避免重复计算滚动对象。
此方案兼顾可读性与性能，是 Pandas 生态中处理「动态窗口滚动计算」的推荐实践。
