Python实现按需求优先级分配库存 替代SQL完成资源分配需求

推荐用pandas实现，性能远优于手写while循环，数千组PARTNER_ID+ITEM_ID的量级下毫秒级就能跑完，核心逻辑是按组计算累计需求后和总供给做比对计算分配量。

实现代码 #

首先导入依赖：

import pandas as pd

数据读取部分可根据实际存储源替换为读csv、读数据库的逻辑，以下为你给出的测试数据示例：

# 构造测试需求表，实际使用替换为pd.read_csv/read_sql等读入逻辑
need_df = pd.DataFrame([
    [1, 'ID32', 621, 57, 1],
    [1, 'ID32', 321, 9, 2],
    [1, 'ID32', 315, 3, 3],
    [1, 'ID32', 732, 1, 4],
    [2, 'ID32', 443, 5, 1],
    [2, 'ID32', 321, 2, 2],
], columns=['PARTNER_ID', 'ITEM_ID', 'STORE', 'NEED', 'NEED_RANK'])

# 构造测试供给表，实际使用替换为对应读入逻辑
supply_df = pd.DataFrame([
    [1, 'ID32', 57],
    [2, 'ID32', 6],
], columns=['PARTNER_ID', 'ITEM_ID', 'SUPPLY'])

核心分配逻辑：

# 1. 需求表按优先级排序，确保同组内排名高的排在前面
need_df = need_df.sort_values(['PARTNER_ID', 'ITEM_ID', 'NEED_RANK'], ignore_index=True)

# 2. 关联总供给量到每一行需求，未匹配到供给的默认赋值为0，可根据业务规则调整
merge_df = need_df.merge(supply_df, on=['PARTNER_ID', 'ITEM_ID'], how='left')
merge_df['SUPPLY'] = merge_df['SUPPLY'].fillna(0)

# 3. 按合作方+商品分组计算累计需求
merge_df['cum_need'] = merge_df.groupby(['PARTNER_ID', 'ITEM_ID'])['NEED'].cumsum()

# 4. 计算每个门店实际分配量
def calc_received(row):
    # 上一行的累计需求 = 当前累计需求 - 当前门店需求
    prev_cum = row['cum_need'] - row['NEED']
    if prev_cum >= row['SUPPLY']:
        return 0
    return min(row['NEED'], row['SUPPLY'] - prev_cum)

merge_df['RECEIVED_SUPPLY'] = merge_df.apply(calc_received, axis=1)

# 5. 输出结果，删除中间计算的辅助字段
result_df = merge_df.drop(['SUPPLY', 'cum_need'], axis=1)

输出的result_df和你给出的目标示例完全一致。

性能优化说明 #

如果后续数据量涨到十万级以上，可以把apply替换为向量化运算，运算速度会提升数倍：

merge_df['prev_cum'] = merge_df['cum_need'] - merge_df['NEED']
merge_df['RECEIVED_SUPPLY'] = 0
# 仅上一轮累计未超过供给的门店可以分配到库存
mask = merge_df['prev_cum'] < merge_df['SUPPLY']
merge_df.loc[mask, 'RECEIVED_SUPPLY'] = pd.concat([
    merge_df.loc[mask, 'NEED'], 
    merge_df.loc[mask, 'SUPPLY'] - merge_df.loc[mask, 'prev_cum']
], axis=1).min(axis=1)

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