优化Pandas计算指定日期前产品平均销售额的高效方案

问题背景 #

现有存储产品销售数据的DataFrame，需要实现函数，根据prod_id、prod_type和指定日期pdate，计算该日期及之前的平均销售额。当前实现每次调用都要全表筛选，数据量大时耗时显著，需优化。

数据构造代码 #

import numpy as np
import pandas as pd
import bisect

product_type = ['A','B']
df = pd.DataFrame({
    'prod_id': np.repeat(np.arange(start=2, stop=5, step=1), 235),
    'prod_type': np.random.choice(np.array(product_type), 705),
    'sales_time': pd.date_range(start='1-1-2018', end='3-30-2018', freq='3H'),
    'sale_amt': np.random.randint(4, 100, size=705)
})

原实现代码（效率较低） #

def cal_avg(product, ptype, pdate):
    temp_df = df[(df['prod_id']==product) & (df['prod_type']==ptype) & (df['sales_time']<= pdate)]
    return temp_df['sale_amt'].mean()

# 调用示例
cal_avg(2,'A','2018-02-12 15:00:00')
# 返回示例结果：53.983

原实现的核心问题是每次调用都要遍历全表做布尔筛选，当数据量达到十万/百万级、且函数调用频繁时，重复的全表扫描会导致严重的性能瓶颈。

优化方案：预处理+二分查找 #

通过一次预处理，将数据按分组排序并预计算累计统计值，后续调用仅需O(logN)的查找操作，大幅提升效率。

步骤1：预处理数据 #

# 1. 按prod_id、prod_type分组，对每组的sales_time排序
grouped = df.groupby(['prod_id', 'prod_type']).apply(
    lambda x: x.sort_values('sales_time').reset_index(drop=True)
).reset_index(drop=True)

# 2. 对每组预计算累计销售额和累计记录数
grouped['cum_sum'] = grouped.groupby(['prod_id', 'prod_type'])['sale_amt'].cumsum()
grouped['cum_count'] = grouped.groupby(['prod_id', 'prod_type']).cumcount() + 1

# 3. 将分组数据转换为字典，方便快速查找
group_dict = {}
for (pid, ptype), g in grouped.groupby(['prod_id', 'prod_type']):
    group_dict[(pid, ptype)] = {
        'times': g['sales_time'].tolist(),
        'cum_sums': g['cum_sum'].tolist(),
        'cum_counts': g['cum_count'].tolist()
    }

步骤2：优化后的计算函数 #

def fast_cal_avg(product, ptype, pdate):
    # 统一转换为datetime类型，避免格式兼容问题
    pdate = pd.to_datetime(pdate)
    # 获取对应分组的数据
    key = (product, ptype)
    if key not in group_dict:
        return np.nan  # 无匹配数据时返回NaN
    
    group_data = group_dict[key]
    # 用二分查找找到第一个大于pdate的位置
    idx = bisect.bisect_right(group_data['times'], pdate)
    
    if idx == 0:
        return np.nan  # 没有早于等于pdate的记录
    # 计算平均：累计总和 / 累计计数
    return group_data['cum_sums'][idx-1] / group_data['cum_counts'][idx-1]

调用示例 #

fast_cal_avg(2,'A','2018-02-12 15:00:00')

优化原理 #

1. 预处理仅执行一次：将数据按分组排序并预计算累计值，避免重复的全表扫描
2. 二分查找替代全表筛选：每次调用通过bisect_right快速定位日期位置，时间复杂度从O(N)降至O(logN)
3. 字典快速索引分组：直接通过(prod_id, prod_type)键获取对应分组数据，无需重复分组计算

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