pandas中含重复时间索引的不规则数据集滚动中位数高效计算方法

原生时间窗口rolling逐行推进计算窗口，遇到重复时间戳的记录时，会依次将每条记录纳入窗口计算，不会等待同时间点的所有记录全部入库后再计算，最终导致同时间点的输出结果逐行跳变，和预期不符。
核心观察：所有索引值完全相同的行，参考时间点一致，对应的窗口范围完全一致，因此最终的滚动结果必然完全相同，不需要逐行重复计算。

利用上述观察，仅在每个时间组的末尾位置用原生rolling计算一次正确结果（此时同时间的所有记录已全部纳入窗口），再将结果广播回同组的所有行，全程使用pandas向量化操作，无逐行Python循环，计算效率较原循环实现提升数个量级。

封装实现代码 #

import numpy as np
import pandas as pd

def fast_irregular_rolling_median(series: pd.Series, window: pd.Timedelta, direction: str = "backward") -> pd.Series:
    """
    高效计算带重复时间索引的不规则滚动中位数
    参数：
        series: 带DatetimeIndex的输入序列，支持重复时间索引，需提前按时间升序排列
        window: 滚动窗口时长
        direction: 滚动方向
            - backward: 向前回溯固定窗口（从当前时间点往过去回溯window时长）
            - forward: 向后延伸固定窗口（从当前时间点往未来延伸window时长）
    """
    if direction == "backward":
        # 原生滚动计算逐行结果
        raw_roll = series.rolling(window).median()
        # 标记每个时间组的末尾位置
        group_end = series.groupby(level=0).cumcount(ascending=False) == 0
        # 仅保留组末尾位置的正确结果，其他位置置空后向后填充
        res = raw_roll.where(group_end).bfill()
        return res
    elif direction == "forward":
        # 倒序序列适配向后窗口的计算逻辑
        rev_series = series.iloc[::-1]
        rev_raw_roll = rev_series.rolling(window).median()
        # 标记倒序序列中每个时间组的末尾位置
        rev_group_end = rev_series.groupby(level=0).cumcount(ascending=False) == 0
        # 倒序下保留正确值、填充、再倒序回原顺序
        rev_res = rev_raw_roll.where(rev_group_end).bfill()
        return rev_res.iloc[::-1]
    else:
        raise ValueError("direction仅支持backward/forward两种取值")

结果验证 #

可直接替换原有逐行循环实现，输出结果与自定义逻辑完全一致：

def dummy_data():
    idx = np.array([pd.Timestamp(year=2021, month=1, day=1, hour=1),
                    pd.Timestamp(year=2021, month=1, day=2, hour=2),
                    pd.Timestamp(year=2021, month=1, day=3, hour=5),
                    pd.Timestamp(year=2021, month=1, day=3, hour=5),
                    pd.Timestamp(year=2021, month=1, day=3, hour=5),
                    pd.Timestamp(year=2021, month=1, day=8, hour=15),
                    pd.Timestamp(year=2021, month=1, day=9, hour=18),
                    pd.Timestamp(year=2021, month=1, day=10, hour=14),
                    ])
    data = np.array([1, 2, 3, 4, 5, 6, 7, 8])
    return pd.DataFrame(data, index=idx, columns=["l"])

def rolling_median_irregular(ds, left, right):
    # 原自定义循环实现
    res = pd.Series(index=ds.index, dtype=np.float64)
    for t in ds.index:
        val = ds.loc[(ds.index >= t - left) & (ds.index <= t + right)].median()
        res.loc[t] = val
    return res

if __name__ == "__main__":
    df = dummy_data()
    # 原自定义计算结果
    df["fun"] = rolling_median_irregular(df["l"], left=pd.Timedelta(days=2), right=pd.Timedelta(days=0))
    df["fun2"] = rolling_median_irregular(df["l"], left=pd.Timedelta(days=0), right=pd.Timedelta(days=2))
    # 新优化方法计算结果
    df["fast_fun"] = fast_irregular_rolling_median(df["l"], window=pd.Timedelta(days=2), direction="backward")
    df["fast_fun2"] = fast_irregular_rolling_median(df["l"], window=pd.Timedelta(days=2), direction="forward")
    print(df)

运行后可观察到：

• 优化后的fast_fun列与原循环计算的fun列结果完全一致，重复时间点的3条记录均得到正确的3.5中位数结果
• 优化后的fast_fun2列与原循环计算的fun2列结果完全一致，重复时间点的3条记录均得到正确的4.0中位数结果
• 全程为向量化计算，无逐行循环，对毫秒级精度的大规模时间序列数据可实现秒级处理

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