如何按列GroupBy后计算满足时间条件的移动均值与累计计数

问题概述 #

对已按country和start_date排序的DataFrame，按country分组后为每行计算两个指标：

• trailing_mean：组内所有end_date早于当前行start_date的历史记录的value均值，无符合条件记录时为NA
• counter：符合上述时间条件的历史记录总数，无符合条件时为0

数据预处理 #

先将日期列转换为datetime类型，确保时间比较的准确性：

import pandas as pd
import numpy as np

# 生成示例数据
data = {'country': ['arg', 'arg', 'arg', 'arg', 'arg', 'usa', 'usa', 'usa'],
        'start_date': ['2020-01-01', '2020-01-01', '2020-05-01', '2021-05-01', '2021-07-01',
                       '2020-03-01', '2020-05-01', '2020-09-01'],
        'end_date': ['2020-10-01', '2020-09-01', '2021-01-01', '2021-06-01', '2021-12-01',
                     '2020-10-01', '2020-08-01', '2021-05-01'],
        'value': [250, 300, 150, 170, 200, 150, 100, 120]}

df = pd.DataFrame(data)

# 转换日期列为datetime类型
df['start_date'] = pd.to_datetime(df['start_date'])
df['end_date'] = pd.to_datetime(df['end_date'])

方法一：直观遍历法（适合小数据集） #

分组后逐行遍历，筛选符合时间条件的记录并计算指标：

def calculate_trailing_metrics(group):
    group['trailing_mean'] = pd.NA
    group['counter'] = 0
    
    for idx, row in group.iterrows():
        # 筛选组内end_date早于当前start_date的记录
        valid_mask = group['end_date'] < row['start_date']
        valid_values = group.loc[valid_mask, 'value']
        
        if len(valid_values) > 0:
            group.loc[idx, 'trailing_mean'] = round(valid_values.mean(), 2)
            group.loc[idx, 'counter'] = len(valid_values)
    
    return group

# 分组应用函数
result_df = df.groupby('country', group_keys=False).apply(calculate_trailing_metrics)

方法二：二分查找优化法（适合大数据集） #

利用numpy的二分查找实现O(n log n)时间复杂度，提升处理效率：

def calculate_trailing_metrics_fast(group):
    # 提取并排序end_date和对应value
    end_dates = group['end_date'].values
    values = group['value'].values
    
    sorted_indices = np.argsort(end_dates)
    sorted_end_dates = end_dates[sorted_indices]
    sorted_values = values[sorted_indices]
    
    # 计算累计和与累计计数
    cumulative_sum = np.cumsum(sorted_values)
    cumulative_count = np.arange(1, len(sorted_values)+1)
    
    # 二分查找找到每个start_date对应的有效记录边界
    positions = np.searchsorted(sorted_end_dates, group['start_date'].values, side='left')
    
    # 赋值结果
    group['counter'] = positions
    group['trailing_mean'] = np.where(positions == 0, pd.NA, round(cumulative_sum[positions-1] / positions, 2))
    
    return group

# 分组应用优化函数
result_df_fast = df.groupby('country', group_keys=False).apply(calculate_trailing_metrics_fast)

输出结果 #

两种方法均会生成符合需求的结果，示例输出如下：

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