优化Pandas箱数分配函数：替代手动iloc与多分支判断

需求回顾 #

对门店员工箱数按以下规则分配：

• 单员工门店：该员工获得全部待分配箱数
• 多员工门店：按员工编号优先级排序，每人最多分配100箱，直到箱数耗尽

原实现通过手动iloc索引和if/elif分支处理，无法适配员工数量≥10的门店，以下是基于循环和分组逻辑的可扩展方案。

完整实现代码 #

import pandas as pd

# 原数据
data_stack_exchange = {'store': ['A','B', 'B', 'C', 'C', 'C', 'D', 'D', 'D', 'D'],
        'worker': [1,1,2,1,2,3,1,2,3,4],
        'boxes': [105, 90, 100, 80, 10, 200, 70, 210, 50, 0],
        'optimal_boxes': [0,0,0,0,0,0,0,0,0,0]}
df_stack_exchange = pd.DataFrame(data_stack_exchange)

def allocate_boxes(group):
    # 按员工编号排序，保证分配优先级
    group_sorted = group.sort_values('worker').reset_index(drop=True)
    # 计算门店总待分配箱数（此处为门店所有员工boxes之和，可根据实际需求调整）
    total_boxes = group_sorted['boxes'].sum()
    remaining_boxes = total_boxes
    allocations = []
    
    # 循环分配，适配任意数量员工
    for _ in range(len(group_sorted)):
        if remaining_boxes <= 0:
            allocations.append(0)
        else:
            # 取剩余箱数和100的较小值作为当前员工分配量
            current_alloc = min(remaining_boxes, 100)
            allocations.append(current_alloc)
            remaining_boxes -= current_alloc
    
    # 将分配结果映射回原分组，并恢复原索引对齐
    group_sorted['optimal_boxes'] = allocations
    return group_sorted.set_index(group.index)

# 按门店分组应用分配逻辑
df_stack_exchange = df_stack_exchange.groupby('store', group_keys=False).apply(allocate_boxes)

print(df_stack_exchange)

方案优势与可扩展性说明 #

1. 适配任意员工数量：通过循环遍历每个员工，无需手动编写分支判断，不管门店有1个还是20个员工都能自动处理
2. 优先级保障：先按员工编号排序，确保分配顺序符合要求
3. 逻辑灵活：总待分配箱数的计算可根据实际需求调整（比如如果boxes列是门店总箱数而非员工个体值，只需将total_boxes改为group_sorted['boxes'].iloc[0]）
4. 索引对齐：处理过程中保留原DataFrame的索引，避免数据错位

输出结果示例 #

执行代码后，optimal_boxes列会得到符合规则的分配结果：

• 门店A（1名员工）：105
• 门店B（2名员工）：100、90
• 门店C（3名员工）：100、100、90
• 门店D（4名员工）：100、100、100、30

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