如何基于Pandas DataFrame订单行条件添加水果处理费行？

需求说明 #

DataFrame中每一行对应一个订单的商品项，若某订单包含至少一个水果类商品（Is_Fruit值为TRUE），则需为该订单追加一行记录：Order为对应订单号，Item为"Handling Fee"（水果处理费），Is_Fruit列留空。

输入示例 #

输出示例 #

现有非优雅实现 #

通过iterrows循环结合跟踪变量实现，代码如下：

import pandas as pd
import numpy as np

# 若订单包含至少一个水果商品，则添加水果处理费
file = r"D:\Dropbox\Python\Tests\fruittest.xlsx"
df = pd.read_excel(file)

# 初始化跟踪变量
prior_order = 0
add_handling = False
df2 = df  # 复制原DataFrame存储结果

for index, row in df.iterrows():
    if index == 0:  # 处理第一行特殊情况
        add_handling = row['Is_Fruit'] 
    
    elif row['Order'] == prior_order:  # 仍在同一订单内
         add_handling = add_handling | row['Is_Fruit']  # 只要有一个是水果，标记为需要添加处理费
        
    elif row['Order'] != prior_order:  # 进入新订单的第一行
        if add_handling:  # 上一订单含水果，添加处理费行
            new_row = {'Order':prior_order, 'Item':'Handling Fee'}
            df2 = df2.append(new_row, ignore_index=True)
            
        add_handling = row['Is_Fruit']  # 重置跟踪标记为新订单的第一个商品状态

    prior_order = row['Order']  # 更新跟踪的订单号

# 循环结束后处理最后一个订单
if add_handling:
    new_row = {'Order':prior_order, 'Item':'Handling Fee'}
    df2 = df2.append(new_row, ignore_index=True)
    
# 按订单号排序结果
df2 = df2.sort_values(by=['Order'], ignore_index=True)

优雅实现方案 #

利用Pandas的分组聚合和合并操作，避免循环，更高效且代码简洁：

实现步骤 #

1. 分组统计每个订单是否包含水果：按Order分组，检查Is_Fruit是否存在TRUE
2. 生成需要添加处理费的订单对应的记录DataFrame
3. 将原DataFrame与处理费记录DataFrame合并，再按Order排序

基础代码实现 #

import pandas as pd

file = r"D:\Dropbox\Python\Tests\fruittest.xlsx"
df = pd.read_excel(file)

# 步骤1：找出所有含水果的订单
fruit_orders = df.groupby('Order')['Is_Fruit'].any().reset_index()
fruit_orders = fruit_orders[fruit_orders['Is_Fruit']]['Order']

# 步骤2：生成处理费行的DataFrame
handling_fee_rows = pd.DataFrame({
    'Order': fruit_orders,
    'Item': 'Handling Fee',
    'Is_Fruit': None  # 留空
})

# 步骤3：合并原数据和处理费行，排序后重置索引
result_df = pd.concat([df, handling_fee_rows], ignore_index=True)
result_df = result_df.sort_values(by=['Order'], ignore_index=True)

print(result_df)

进阶：确保处理费行在订单末尾 #

如果需要严格保证处理费行出现在对应订单的所有商品之后，可以添加排序键优化：

import pandas as pd

file = r"D:\Dropbox\Python\Tests\fruittest.xlsx"
df = pd.read_excel(file)

# 给原数据添加订单内的排序序号
df['sort_idx'] = df.groupby('Order').cumcount()

# 找出含水果的订单并生成处理费行，排序序号设为对应订单商品数+1
fruit_orders = df.groupby('Order')['Is_Fruit'].any()
handling_fee_rows = pd.DataFrame({
    'Order': fruit_orders[fruit_orders].index,
    'Item': 'Handling Fee',
    'Is_Fruit': None,
    'sort_idx': df.groupby('Order')['sort_idx'].max() + 1
})

# 合并后按订单号+排序序号排序，最后删除辅助列
result_df = pd.concat([df, handling_fee_rows], ignore_index=True)
result_df = result_df.sort_values(by=['Order', 'sort_idx'], ignore_index=True).drop('sort_idx', axis=1)

print(result_df)

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