基于外部函数的Pandas双DataFrame任务-人员分配方案优化咨询

需求说明 #

现有两个pandas DataFrame：

• 一个存储员工列表，包含员工ID和姓名
• 一个存储任务列表，包含任务ID、日期、班次、所需员工数

需要为每个任务分配员工，规则是：逐个处理任务，对每个任务遍历员工，通过外部函数返回1（分配）或0（跳过）判断是否分配该员工，直到任务凑够所需人数，再处理下一个任务。当前用随机数模拟外部函数的决策结果，原代码存在循环控制冗余问题，需优化遍历逻辑。

数据结构 #

任务DataFrame #

task_id  day  shift  workers_needed
0       1    4      1               3
1       2    4      2               5
2       3    7      1               4
3       4    5      2               4

员工DataFrame #

worker_id        name
0        101  Worker_101
1        102  Worker_102
2        103  Worker_103
3        104  Worker_104

原代码问题 #

原代码的循环控制逻辑冗余：

• 提前初始化任务变量，在循环内重复更新
• 边界判断繁琐，容易出错
• 未清晰区分任务遍历和员工分配的逻辑

优化后的实现方案 #

核心思路 #

1. 用iterrows()遍历任务DataFrame，直接获取每个任务的信息，避免手动索引控制
2. 封装外部判断函数，明确职责
3. 对每个任务，循环尝试分配员工直到满足需求，逻辑清晰易维护

优化代码 #

import pandas as pd
import random

# 模拟外部判断函数：返回1表示分配，0表示跳过
def should_assign_worker():
    return random.randint(0, 1)

# 示例数据
tasks_data = {
    'task_id': [1, 2, 3, 4],
    'day': [4, 4, 7, 5],
    'shift': [1, 2, 1, 2],
    'workers_needed': [3, 5, 4, 4]
}
workers_data = {
    'worker_id': [101, 102, 103, 104],
    'name': ['Worker_101', 'Worker_102', 'Worker_103', 'Worker_104']
}

tasks_df = pd.DataFrame(tasks_data)
workers_df = pd.DataFrame(workers_data)

# 存储分配结果
assignments = []

# 遍历每个任务
for _, task in tasks_df.iterrows():
    task_id = task['task_id']
    needed = task['workers_needed']
    assigned_count = 0
    assigned_workers = []
    
    print(f"开始处理任务 {task_id}，需要 {needed} 名员工")
    
    # 遍历员工，直到凑够所需人数
    for _, worker in workers_df.iterrows():
        if assigned_count >= needed:
            break
        
        if should_assign_worker() == 1:
            assigned_count += 1
            assigned_workers.append(worker['worker_id'])
            print(f"  分配员工 {worker['worker_id']}，已分配 {assigned_count}/{needed}")
    
    # 记录该任务的分配结果
    assignments.append({
        'task_id': task_id,
        'assigned_workers': assigned_workers,
        'total_assigned': assigned_count
    })

# 输出最终分配结果
print("\n最终分配结果：")
for res in assignments:
    print(f"任务 {res['task_id']}：分配了 {res['total_assigned']} 名员工，ID为 {res['assigned_workers']}")

代码说明 #

• iterrows()遍历：直接迭代DataFrame的行，无需手动维护索引变量i，逻辑更直观
• 封装判断函数：should_assign_worker()单独处理分配判断，便于后续替换为真实业务逻辑
• 清晰的分层逻辑：先遍历任务，再对每个任务遍历员工，直到满足需求，流程清晰
• 结果存储：用列表记录每个任务的分配结果，方便后续分析或存储

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