基于Pandas的序贯分配问题高效求解方案问询

问题背景 #

需要将物品供应数据映射到对应消耗周次及数量，满足特定分配规则，现有基于循环的方案在大规模数据下效率不足，寻求优化方案。

数据示例 #

供应数据（supply_df） #

supply_df = pd.DataFrame({
    'item': ['a']*9 + ['b']*9,
    'supplier': ['s1','s2','s3','s1','s2','s3','s1','s2','s3']*2,
    'supply_week': [1,1,1,2,2,2,3,3,3]*2,
    'supply': [10,20,10,23,33,42,52,27,29,37,32,38,17,28,44,41,45,24]
})

消耗数据（consume_df） #

consume_df = pd.DataFrame({
    'item': ['a','a','a','b','b','b'],
    'week': [1,2,3,1,2,3],
    'consume': [33,100,102,90,100,80]
})

分配规则 #

• 同周供应优先按s1 > s2 > s3的顺序消耗
• 遵循先到先消耗：早周次供应优先被消耗
• 供应仅能用于对应物品的消耗
• 单周消耗可来自多供应商，单供应商供应可拆分到多周
• 剩余供应无需考虑

现有方案痛点 #

使用pd.DataFrame.GroupBy.apply结合itertuples/iterrows循环处理，在数据量较大时，循环操作带来的性能瓶颈明显，效率低下。

高效解决方案 #

采用矢量化操作+区间匹配的思路，避免逐行循环，利用pandas的高效内置函数提升处理速度：

步骤1：预处理供应数据 #

按物品分组，按供应周次升序、供应商优先级排序，计算累计供应量：

import pandas as pd

# 定义供应商优先级映射（数值越小优先级越高）
supplier_priority = {'s1': 1, 's2': 2, 's3': 3}

supply_processed = (
    supply_df
    .assign(priority=lambda x: x['supplier'].map(supplier_priority))
    .sort_values(by=['item', 'supply_week', 'priority'])
    .groupby('item', group_keys=False)
    .assign(cum_supply=lambda x: x['supply'].cumsum())
    .reset_index(drop=True)
)

步骤2：预处理消耗数据 #

按物品分组，按周次升序，计算累计消耗量：

consume_processed = (
    consume_df
    .sort_values(by=['item', 'week'])
    .groupby('item', group_keys=False)
    .assign(cum_consume=lambda x: x['consume'].cumsum())
    .reset_index(drop=True)
)

步骤3：区间匹配分配供应 #

利用merge_asof高效匹配累计供应与累计消耗的区间，计算每个供应对应的消耗周次及数量：

result_list = []
for item in supply_processed['item'].unique():
    # 提取当前物品的供应和消耗数据
    supply_item = supply_processed[supply_processed['item'] == item].copy()
    consume_item = consume_processed[consume_processed['item'] == item].copy()
    
    # 添加初始累计消耗为0的行，处理第一个供应的匹配
    consume_item = pd.concat([
        pd.DataFrame({'item': [item], 'week': [0], 'consume': [0], 'cum_consume': [0]}),
        consume_item
    ], ignore_index=True)
    
    # 用merge_asof找到每个累计供应对应的首个覆盖的累计消耗区间
    merged = pd.merge_asof(
        supply_item.sort_values('cum_supply'),
        consume_item.sort_values('cum_consume'),
        left_on='cum_supply',
        right_on='cum_consume',
        direction='forward',
        by='item'
    )
    
    # 计算当前供应在消耗区间内的分配数量
    merged['prev_cum_consume'] = merged['cum_consume'].shift(1).fillna(0)
    merged['start'] = merged[['cum_supply', 'prev_cum_consume']].max(axis=1)
    merged['end'] = merged[['cum_supply', 'cum_consume']].min(axis=1)
    merged['allocated'] = merged['end'] - merged['start']
    
    # 过滤无分配的记录，关联消耗周次信息
    merged = merged[merged['allocated'] > 0].merge(
        consume_item[['item', 'week', 'cum_consume']],
        on=['item', 'cum_consume'],
        suffixes=('', '_consume')
    )
    
    # 整理结果列
    result_list.append(merged[['item', 'supplier', 'supply_week', 'week', 'allocated']])

# 合并最终结果
out_df = pd.concat(result_list, ignore_index=True)

方案优势 #

• 全程使用pandas矢量化操作，避免逐行循环，处理大规模数据时性能提升显著
• merge_asof是专门针对有序数据的高效区间匹配函数，时间复杂度远低于循环遍历
• 分组处理逻辑清晰，兼顾可读性与性能

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