如何基于DataFrame分组行创建位置分配计算列？

输入数据 #

import pandas as pd

df = pd.DataFrame({'ID': ['ID1', 'ID1', 'ID2', 'ID3', 'ID4', 'ID4'],
                   'CONTAINER': ['CAT1', 'CAT2', 'CAT3', 'CAT2', 'CAT1', 'CAT2'],
                   'FREE_POSITIONS': [1, 5, 32, 0, 5, 10]})

needed_positions = {"ID1": 3, "ID2": 5, "ID3": 2, "ID4": 2}

需求说明 #

需要新增POSITIONS_TO_ALLOCATE列，逻辑为：

• 按ID分组，每个ID对应总需求（来自needed_positions字典）
• 组内按行顺序分配：优先用当前行的空闲位置满足剩余需求，分配后更新剩余需求
• 剩余需求为0时，后续行分配0

正确实现代码 #

def allocate_positions(group):
    # 获取当前分组ID对应的总需求
    remaining_need = needed_positions[group['ID'].iloc[0]]
    alloc_list = []
    for free in group['FREE_POSITIONS']:
        if remaining_need <= 0:
            alloc_list.append(0)
            continue
        # 取空闲位置和剩余需求的较小值作为分配量
        alloc = min(free, remaining_need)
        alloc_list.append(alloc)
        remaining_need -= alloc
    group['POSITIONS_TO_ALLOCATE'] = alloc_list
    return group

# 按ID分组处理，生成结果
result = df.groupby('ID', group_keys=False).apply(allocate_positions)
print(result)

运行结果 #

ID CONTAINER  FREE_POSITIONS  POSITIONS_TO_ALLOCATE
0  ID1      CAT1               1                      1
1  ID1      CAT2               5                      2
2  ID2      CAT3              32                      5
3  ID3      CAT2               0                      0
4  ID4      CAT1               5                      2
5  ID4      CAT2              10                      0

原代码问题分析 #

1. 未跟踪剩余需求：每次直接拿总需求和当前行空闲位置比较，无法处理需求递减的情况
2. 错误添加分组：循环内每次append整个分组，会导致数据重复冗余
3. 缺少else分支逻辑：未处理空闲位置小于总需求时的分配逻辑

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