基于df2的count值快速筛选df1对应行数的Pandas优化方法问询

针对你用iterrows循环导致速度过慢的问题，推荐两种矢量化的高效实现方案，底层基于Pandas的C优化操作，能大幅提升处理速度：

方案一：合并后用cumcount筛选 #

通过合并获取每个分组的提取行数，再计算组内行号进行筛选：

# 合并df1与df2，将每个分组的count值关联到原始数据
df_merged = df1.merge(df2, on=['site_id', 'date', 'hour'], how='left')

# 按分组计算组内从0开始的行号
df_merged['row_num'] = df_merged.groupby(['site_id', 'date', 'hour']).cumcount()

# 筛选行号小于对应count的记录，最后清理临时字段
result = df_merged[df_merged['row_num'] < df_merged['count']].drop(columns=['row_num', 'count'])

方案二：分组映射+head()提取 #

先构建分组到提取行数的映射，再分组提取前N行：

# 将df2转换为分组元组到count值的字典
count_map = df2.set_index(['site_id', 'date', 'hour'])['count'].to_dict()

# 按分组提取对应行数，group_keys=False避免保留分组索引
result = df1.groupby(['site_id', 'date', 'hour'], group_keys=False).apply(
    lambda x: x.head(count_map.get(x.name, 0))
)

效率说明 #

这两种方案都避免了纯Python循环，利用Pandas的矢量化分组、合并操作，底层由C实现，在数据量较大时，效率比iterrows高几十甚至上百倍。

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