Pandas大数据量下按id分组生成A_flag列的性能优化问题

原方案性能瓶颈来自两个部分：一是itertuples逐行遍历属于Python级循环，开销远高于pandas向量化操作；二是额外创建临时列增加了内存开销。可选用以下两种优化方案，性能均可提升数倍至数十倍，超大数据量优先选第二种。

方案1：简洁分组写法（中小数据量适用） #

直接通过分组+transform一次性完成计算，无需创建临时列，代码更简洁：

import pandas as pd

# 计算A_flag
df['A_flag'] = df.groupby('id')['type']\
                 .transform(lambda x: x.eq('A').any())\
                 .astype('int')

逻辑说明：按id分组后判断每组内是否存在type等于A的记录，返回布尔值后直接转为0/1即可。

方案2：集合匹配写法（超大数据量最优） #

完全避免分组聚合开销，仅需两次向量化操作，千万行级数据也可快速处理：

# 先提取所有存在type=A的id，存入集合
a_id_set = set(df.loc[df['type'] == 'A', 'id'])
# 直接判断每行的id是否在上述集合中，转成int即得目标列
df['A_flag'] = df['id'].isin(a_id_set).astype('int')

逻辑说明：先筛选出所有type为A的id去重，再直接匹配每行id是否属于该集合，时间复杂度为O(n)，远低于分组聚合的开销。

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