Pandas多DataFrame关联计算优化：移除循环提升运行效率

看起来你现在用嵌套循环处理这个统计需求，数据量大的时候肯定卡得不行——毕竟Python循环在处理DataFrame时效率真的很低。我给你一套完全矢量化的pandas方案，速度能提升几个数量级，而且代码更简洁易维护。

先回顾下你的数据结构，方便大家理解：

原始DataFrame构造代码 #

import pandas as pd
import numpy as np

# df_main构造
users = ['id1','id1','id2','id2','id3','id3','id4']
keywords = ['k1','k1', 'k2','k2','k2','k3','k3']
quantity = [10,10,2,2,2,4,4]
duration = [1,1,3,3,3,2,2]
df_main = pd.DataFrame(list(zip(users, keywords, quantity, duration)), columns = ['users','keywords','quantity','duration'])

# df1构造
columns = ['USER_ID','k1','k2','k3']
users = ['id1','id2','id3','id4']
values1 = [1,0,0,0]
values2 = [0,1,1,0]
values3 = [0,0,1,1]
df1 = pd.DataFrame(list(zip(users, values1, values2, values3)), columns = columns)

优化方案思路 #

核心是利用pandas的分组聚合、宽表转长表（melt）、**长表转宽表（pivot）**这些内置矢量化操作，完全避开Python循环：

• 先聚合明细数据：把df_main里每个用户-关键词组合的quantity和duration总和一步算到位
• 转换标记表为长格式：把df1的宽表结构转成每行对应一个用户-关键词有效对（标记为1的），方便后续合并
• 合并+重塑结果：把聚合数据和标记数据合并后，转回宽表结构对齐原df1格式，最后补全0值

完整优化代码 #

# 步骤1：聚合df_main的用户-关键词统计量
agg_df = df_main.groupby(['users', 'keywords']).agg(
    total_quantity=('quantity', 'sum'),
    total_duration=('duration', 'sum')
).reset_index()

# 步骤2：将df1转成长格式，筛选出标记为1的用户-关键词对
melted_df1 = df1.melt(
    id_vars='USER_ID',
    value_vars=['k1', 'k2', 'k3'],
    var_name='keywords',
    value_name='flag'
).query('flag == 1').drop('flag', axis=1)

# 步骤3：合并聚合数据与标记数据，再转回宽表
merged = melted_df1.merge(agg_df, left_on=['USER_ID', 'keywords'], right_on=['users', 'keywords'], how='left')

# 重塑为宽表，调整列名格式
wide_result = merged.pivot(
    index='USER_ID',
    columns='keywords',
    values=['total_quantity', 'total_duration']
).reset_index()

# 重命名列，比如把('total_quantity', 'k1')改成k1_quantity
wide_result.columns = [
    f'{col[1]}_{col[0].replace("total_", "")}' 
    if col[0] != 'USER_ID' 
    else col[0] 
    for col in wide_result.columns
]

# 合并回原df1，补全标记为0的关键词对应的统计值为0
final_result = df1.merge(wide_result, on='USER_ID', how='left').fillna(0)

# 查看结果
print(final_result)

为什么这个方案更快？ #

原循环是嵌套遍历关键词和用户，每次都要做布尔索引和求和——这相当于每次循环都要扫描一遍整个DataFrame，时间复杂度是O(n*m)（n是关键词数量，m是用户数量）。

而优化后的方案用的是pandas底层实现的矢量化操作，这些操作都是用C语言编写的，不需要Python解释器逐行执行，时间复杂度接近O(n)，数据量越大，性能提升越明显。

最终输出结果 #

运行后你会得到完全符合需求的结果：

USER_ID  k1  k2  k3  k1_quantity  k1_duration  k2_quantity  k2_duration  k3_quantity  k3_duration
0     id1   1   0   0         20.0          2.0          0.0          0.0          0.0          0.0
1     id2   0   1   0          0.0          0.0          4.0          6.0          0.0          0.0
2     id3   0   1   1          0.0          0.0          2.0          3.0          4.0          2.0
3     id4   0   0   1          0.0          0.0          0.0          0.0          4.0          2.0

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