如何使用DataFrame另一列值作为动态窗口计算pandas rolling滚动和

核心思路 #

pandas原生rolling方法默认只支持固定窗口大小，要实现动态窗口（每行窗口大小由span_pre列指定），可以通过自定义窗口索引器或者分组逐行计算两种方案实现，两种方法都不需要显式编写大循环。

方案1：自定义窗口索引器（适合大数据量，性能更高） #

基于pandas提供的BaseIndexer自定义动态窗口的边界规则，配合分组操作避免不同ID的数据互相干扰：

步骤1：导入依赖&构造示例数据 #

import pandas as pd
import numpy as np
from pandas.api.indexers import BaseIndexer

# 示例数据构造
data = {
    'ID': ['A']*6 + ['S']*6,
    'span_pre': [0,1,2,3,3,3,0,1,2,3,3,3],
    'value1': [0.1,0.0,0.1,0.1,0.1,0.1,0.2,0.2,0.2,np.nan,np.nan,np.nan],
    'value2': [0]*6 + [0]*3 + [1]*3
}
df = pd.DataFrame(data)

步骤2：定义动态窗口索引器 #

class DynamicWindow(BaseIndexer):
    def get_window_bounds(self, values, window, min_periods, center, closed, step):
        n = len(values)
        span_pre = self.extra_params["span_pre"]
        # 窗口逻辑可根据实际需求调整，示例逻辑：窗口长度为span_pre+1，不足时从组内第一行开始取
        # 若需最大窗口限制，可修改为 window_len = np.minimum(span_pre + 1, 3)
        window_len = span_pre + 1
        start = np.maximum(0, np.arange(n) - (window_len -1))
        end = np.arange(n) + 1
        return start, end

步骤3：分组计算滚动和 #

def calc_group_sum(group):
    indexer = DynamicWindow(extra_params={"span_pre": group["span_pre"].values})
    # min_periods=1 表示窗口内只要有1个有效值就返回和，全空返回NaN
    group["SUM_pre"] = group["value1"].rolling(window=indexer, min_periods=1).sum()
    return group

df = df.groupby("ID", group_keys=False).apply(calc_group_sum)

方案2：分组逐行apply（适合小数据量，逻辑更灵活易调整） #

如果数据量不大，可以直接在分组内逐行计算窗口和，调整窗口边界更方便：

def row_sum(row, group_df):
    # 可根据实际窗口规则调整start、end的计算逻辑
    window_len = row["span_pre"] + 1
    end_pos = row.name + 1
    start_pos = max(0, end_pos - window_len)
    # min_count=1 匹配全空窗口返回NaN的需求
    return group_df.loc[start_pos:end_pos-1, "value1"].sum(min_count=1)

df["SUM_pre"] = df.groupby("ID", group_keys=False).apply(
    lambda g: g.apply(row_sum, group_df=g, axis=1)
)

结果验证 #

运行后调整窗口逻辑即可得到和预期完全一致的SUM_pre列，若你实际的窗口规则和示例有差异，只需要修改start和end的计算规则即可。

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