Pandas遍历ID子集列耗时过长，求高效替代方案

你的代码瓶颈在于循环遍历每个ID做子集筛选——对于50万个唯一ID，每次df[df.IDs == id]都会全表扫描一次，累计下来的时间开销是灾难性的。直接用Pandas的多层分组+分组内计算可以彻底解决这个问题，完全避免循环。

优化思路 #

1. 先一次性按IDs和Month做分组，计算column1的均值，这一步是Pandas原生的向量化操作，效率远高于循环筛选后再分组。
2. 再按IDs分组，对每个ID下的月度均值执行EWM求和，同样用向量化的分组操作替代循环。

优化后的代码 #

import pandas as pd

df = pd.DataFrame({"IDs": [1, 1, 1, 2, 2, 2, 3, 3, 3],
                   "Month": ["01", "02", "01", "01", "02", "01", "01", "02", "01"],
                   "column1": [0.9, 0.5, 0.3, 0.8, 0.5, 0.1, 0.6, 0.2, 0.8]})

# 第一步：一次性计算每个ID+Month的均值
grouped_mean = df.groupby(["IDs", "Month"])["column1"].mean().reset_index(level=1)

# 第二步：按ID分组，执行EWM求和
result = grouped_mean.groupby("IDs")["column1"].apply(
    lambda x: x.ewm(span=3, adjust=True).sum()
)

# 如果需要和原代码一样的列表格式，可以转成列表
df_list = [group for _, group in result.groupby("IDs")]

性能说明 #

• 原代码的时间复杂度是O(M*N)（M是唯一ID数，N是总记录数），优化后的代码是O(N)，对于600万条记录+50万唯一ID的场景，耗时会从十几小时压缩到几分钟级别。
• 核心是利用Pandas的groupby向量化操作，避免了循环中的重复数据切片和全表扫描。

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