如何在Pandas中实现带条件的滚动均值计算？

给定已按start_time排序的DataFrame，需为每行添加rolling_mean列——该值为当前行之前所有end_time ≤ 当前行start_time的历史行的duration均值。由于数据量过大，自合并过滤的方法性能不足，以下是O(n log n)复杂度的高效实现方案，基于**树状数组（Fenwick Tree）**实现：

实现步骤 #

1. 离散化数值：将所有end_time和start_time的唯一值映射为整数索引，用于树状数组的操作。
2. 树状数组维护：按DataFrame的行顺序（即start_time从小到大）遍历，先查询树状数组中符合条件的历史行的duration总和与数量，计算均值；再将当前行的end_time和duration插入树状数组，供后续行查询使用。

完整代码 #

import pandas as pd
import numpy as np

# 示例数据
d = [[1, 5],
[2, 5],
[3, 4],
[4, 6],
[5, 9],
[6, 7],
[7, 10]]
df = pd.DataFrame(d, columns=['start_time', 'end_time'])
df['duration'] = df.end_time - df.start_time
# 确保DataFrame按start_time排序（若未排序则执行此步）
df = df.sort_values('start_time').reset_index(drop=True)

# 定义树状数组类，支持点更新和区间查询（sum与count）
class FenwickTree:
    def __init__(self, size):
        self.size = size
        self.sum_tree = [0] * (size + 1)
        self.count_tree = [0] * (size + 1)
    
    def update(self, idx, value):
        # 从idx开始向上更新树状数组
        while idx <= self.size:
            self.sum_tree[idx] += value
            self.count_tree[idx] += 1
            idx += idx & -idx
    
    def query(self, idx):
        # 查询<=idx的所有元素的sum和count
        sum_val = 0
        count_val = 0
        while idx > 0:
            sum_val += self.sum_tree[idx]
            count_val += self.count_tree[idx]
            idx -= idx & -idx
        return sum_val, count_val

# 离散化所有涉及的时间值（end_time和start_time）
all_time_values = np.concatenate([df['end_time'], df['start_time']])
unique_times = np.unique(all_time_values)
time_to_idx = {time: idx + 1 for idx, time in enumerate(unique_times)}  # 索引从1开始

# 初始化树状数组
fenwick_tree = FenwickTree(len(unique_times))

# 遍历计算rolling_mean
rolling_means = []
for start, end, dur in zip(df['start_time'], df['end_time'], df['duration']):
    # 查询符合条件的历史行的总和与数量
    total_dur, row_count = fenwick_tree.query(time_to_idx[start])
    if row_count > 0:
        rolling_means.append(total_dur / row_count)
    else:
        rolling_means.append(np.nan)
    # 将当前行的end_time和duration插入树状数组
    fenwick_tree.update(time_to_idx[end], dur)

df['rolling_mean'] = rolling_means

结果验证 #

运行代码后，示例DataFrame的rolling_mean列结果如下：

性能说明 #

• 时间复杂度：O(n log m)，其中n是DataFrame行数，m是离散化后的唯一时间值数量，远优于自合并的O(n²)，可轻松处理百万级数据。
• 空间复杂度：O(m)，仅需存储离散化映射和树状数组，内存占用可控。

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