Pandas优化问询：合并DataFrame并按占比拆分分配订单量

现有两个Pandas DataFrame如下：

import pandas as pd

proj_df = pd.DataFrame({'reg_id':[1,2,3,4,5,6,7],
                              'partner': ['ABC_123','ABC_123','ABC_123','ABC_123','ABC_123','ABC_123','ABC_123'],
                              'part_no':['P123','P123','P123','P123','P123','P123','P123'],
                              'cust_info':['Apple','Apple','Apple','Apple','Apple','Apple','Tesla'],
                              'qty_1st_year':[100,100,600,150,50,0,10]})

order_df = pd.DataFrame({'partner': ['ABC_123','ABC_123','JKL_123','MNO_123'],
                         'part_no':['P123','P123','Q123','P567'],
                         'cust_info':['Apple','Hyundai','REON','Renault'],
                         'order_qty':[1000,600,50,0]})

需求说明 #

• 基于partner、part_no、cust_info三个字段合并上述两个DataFrame；
• 将order_df中的order_qty字段拆分，按比例分配至新列assigned_qty；
• 分配比例规则：每组（partner+part_no+cust_info）内，单个qty_1st_year值除以该组qty_1st_year的总和，得到占比后乘以对应order_qty。

当前实现方案 #

已编写如下代码实现需求，但寻求更优雅高效的优化方案：

sum_df = proj_df.groupby(['partner','part_no','cust_info'])['qty_1st_year'].sum().reset_index()
sum_df.columns = ['partner','part_no','cust_info','total_qty_all_project']

t1=proj_df.merge(order_df,on=['partner','part_no','cust_info'],how='left')
t2 = t1.merge(sum_df,on=['partner','part_no','cust_info'],how='left')
t2['pct_value'] = (t2['qty_1st_year']/t2['total_qty_all_project'])*100
proj_df['assigned_value'] = (t2['order_qty']*t2['pct_value'])/100

预期输出 #

预期输出为在原proj_df基础上新增assigned_qty列，具体结果如下：

优化方案 #

可以通过groupby.transform直接在合并后的DataFrame上计算组内总和，减少不必要的merge操作，同时简化计算逻辑，代码如下：

# 合并两个DataFrame，保留proj_df的所有行
merged_df = proj_df.merge(order_df, on=['partner', 'part_no', 'cust_info'], how='left')

# 计算每组的qty_1st_year总和，直接生成新列
merged_df['total_group_qty'] = merged_df.groupby(['partner', 'part_no', 'cust_info'])['qty_1st_year'].transform('sum')

# 计算分配量，处理组内总和为0的情况避免除零错误
merged_df['assigned_qty'] = merged_df.apply(
    lambda row: (row['qty_1st_year'] / row['total_group_qty']) * row['order_qty'] 
    if row['total_group_qty'] != 0 else 0,
    axis=1
)

# 将结果同步回原proj_df
proj_df['assigned_qty'] = merged_df['assigned_qty']

优化点说明 #

• 使用transform替代单独分组求和后merge，减少一次数据合并操作，提升效率；
• 直接计算分配量，去掉中间百分比计算步骤，简化逻辑；
• 增加除零判断，避免组内qty_1st_year总和为0时出现运算错误。

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