如何用groupby高效计算DataFrame中前N个月的销售总额

嘿，这个问题我太熟了——你现在用的列表推导式逐行查询的方式，在数据量小的时候还能凑合用，但数据一上来就会拖慢速度，毕竟每一行都要全表扫描一次，完全是重复劳动！咱们换个思路，用groupby先把月度销售总额预处理好，再映射回原DataFrame，效率能提升一大截。

问题根源分析 #

你写的列表推导式：

df_test['prev_1m_sales'] = [ sum(df_test.loc[df_test['ym'] == x].sales) for x in df_test['prev_1_ym'] ]

本质上是循环遍历prev_1_ym的每一个值，每次都执行一次df.loc的筛选求和。假设你的DataFrame有1万行，那就要做1万次全表扫描，时间复杂度是O(n²)，数据量越大越慢。

优化方案：先聚合再映射 #

我们只需要先做一次groupby聚合，得到每个月份的销售总额，然后用map方法把往期月份对应的总额直接映射回去，整个过程只需要O(n)的时间复杂度。

步骤如下：

1. 计算月度销售总额：用groupby按ym分组，对sales求和，得到一个以ym为索引的Series。
2. 映射往期销售总额：用这个Series的map方法，把prev_1_ym和prev_2_ym对应的月度总额直接提取到原DataFrame中。

具体代码实现

import pandas as pd

# 先构建你的测试数据（和你给出的一致）
date = ['28-01-2017','29-01-2017','30-01-2017','31-01-2017','01-02-2017','02-02-2017']
sales = [1,2,3,4,1,2]
ym = [201701,201701,201701,201701,201702,201702]
prev_1_ym = [201612,201612,201612,201612,201701,201701]
prev_2_ym = [201611,201611,201611,201611,201612,201612]
df_test = pd.DataFrame({'date': date,'ym':ym,'prev_1_ym':prev_1_ym,'prev_2_ym':prev_2_ym,'sales':sales})
df_test['date'] = pd.to_datetime(df_test['date'],format = '%d-%m-%Y')

# 核心优化步骤：
# 1. 按月份分组计算销售总额
monthly_sales = df_test.groupby('ym')['sales'].sum()

# 2. 映射前1个月和前2个月的销售总额
df_test['prev_1m_sales'] = df_test['prev_1_ym'].map(monthly_sales)
df_test['prev_2m_sales'] = df_test['prev_2_ym'].map(monthly_sales)

# 如果存在往期月份没有销售数据的情况（比如示例中的201612、201611），可以用fillna填充为0
df_test['prev_1m_sales'] = df_test['prev_1m_sales'].fillna(0)
df_test['prev_2m_sales'] = df_test['prev_2m_sales'].fillna(0)

print(df_test)

为什么这个方法更快？ #

• groupby只需要对全表做一次聚合操作，一次性计算出所有月份的销售总额。
• map操作是基于索引的快速查找，相当于字典的键值对查询，速度极快。
• 整体时间复杂度从原来的O(n²)降到了O(n)，数据量越大，提速效果越明显。

输出结果验证 #

运行上面的代码后，你会得到和原来列表推导式一致的结果（缺失的往期月份会被填充为0）：

date      ym  prev_1_ym  prev_2_ym  sales  prev_1m_sales  prev_2m_sales
0 2017-01-28  201701     201612     201611      1            0.0            0.0
1 2017-01-29  201701     201612     201611      2            0.0            0.0
2 2017-01-30  201701     201612     201611      3            0.0            0.0
3 2017-01-31  201701     201612     201611      4            0.0            0.0
4 2017-02-01  201702     201701     201612      1           10.0            0.0
5 2017-02-02  201702     201701     201612      2           10.0            0.0

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