如何实现Pandas中按反向累加阈值定义的自定义滚动窗口

需求说明 #

现有包含A、B两列的DataFrame，需要从每行出发向前累加A列的元素，直到累加和大于等于指定阈值，生成对应的行索引窗口，再基于窗口对B列做自定义聚合计算，最终支持类似df.my_rolling(on='A', func='sum', threshold=10).B.mean()的链式调用效果。

实现步骤 #

1. 构造测试数据 #

import pandas as pd
import numpy as np

df = pd.DataFrame({
    'A': [3,2,4,4,1,7,5,3],
    'B': [10,20,30,40,50,60,70,80] # 示例B列值方便验证计算结果
})

2. 核心实现 #

方案1：基于时间滚动窗口的标准方法（匹配你提到的转换datetime思路）

利用pandas原生支持的时间滚动窗口特性，把A列累加和转换为时间差作为索引，直接用标准rolling方法匹配阈值：

threshold = 10
# 计算A列累计和，转换为秒级时间差作为索引
df['cumsum_A'] = df['A'].cumsum()
df = df.set_index(pd.to_timedelta(df['cumsum_A'], unit='s'))

# 定义阈值大小的时间窗口，closed='right'表示包含当前行
rolling_obj = df.rolling(window=f'{threshold}s', closed='right')

# 生成窗口索引列，索引顺序和你示例完全一致
df['window_indices'] = rolling_obj.apply(lambda x: list(reversed(x.index.astype(int).tolist())))['A']
# 直接对B列做聚合计算
df['B_rolling_mean'] = rolling_obj['B'].mean()

# 过滤累加和不足阈值的空窗口行
result = df[df['cumsum_A'] >= threshold].reset_index(drop=True)

方案2：自定义扩展访问器，实现链式调用效果

通过pandas扩展能力封装自定义滚动逻辑，完全匹配你需要的调用方式：

@pd.api.extensions.register_dataframe_accessor("my_rolling")
class ThresholdRolling:
    def __init__(self, obj):
        self._obj = obj
    
    def __call__(self, on, threshold):
        df = self._obj.copy()
        cumsum_val = df[on].cumsum()
        window_list = []
        valid_mask = []
        for i in range(len(df)):
            # 二分查找窗口左边界，保证区间A列累加和>=阈值
            left_bound = np.searchsorted(cumsum_val, cumsum_val[i] - threshold, side='right')
            current_sum = cumsum_val[i] - (cumsum_val[left_bound-1] if left_bound>0 else 0)
            if current_sum >= threshold:
                # 窗口索引反转，和你示例输出顺序一致
                window_list.append(list(reversed(range(left_bound, i+1))))
                valid_mask.append(True)
            else:
                window_list.append(pd.NA)
                valid_mask.append(False)
        df['window_indices'] = window_list
        # 只返回满足阈值要求的行
        return df[valid_mask].copy()

调用示例：

# 完全匹配你示例的调用方式
mean_b = df.my_rolling(on='A', threshold=10).B.mean()
print(mean_b)

结果验证 #

实现效果和你给出的预期窗口完全一致，索引3对应窗口[3,2,1]，A列累加和4+4+2=10刚好满足阈值要求。

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