迭代查找订单可用产能日期及大数量级脚本优化问询

现有数据集 #

import pandas as pd

sales_data = {
    'order_number': [1001, 1002, 1003, 1004, 1005],
    'order_date': ['2022-02-01', '2022-02-03', '2022-02-07', '2022-02-10', '2022-02-14']
}
sales_data = pd.DataFrame(sales_data)

capacity_data = {
    'date': pd.date_range(start='2022-02-01', end='2022-02-28', freq='D'),
    'capacity': [0, 0, 0, 1, 1, 100, 110, 120, 130, 140, 150, 160, 170, 180, 
                 190, 200, 210, 220, 230, 240, 250, 260, 270, 280, 290, 300, 310]
}
capacity_data = pd.DataFrame(capacity_data)

期望输出 #

output = {
    'order_number': [1001, 1002, 1003, 1004, 1005], 
    'confirmation_date':['2022-02-04', '2022-02-05', '2022-02-06', '2022-02-06', '2022-02-06']
}

需求说明 #

为每个订单查找最近的可用产能日期（capacity>0），这里的“最近”指晚于订单日期的最早可用日期，占用该产能后将对应日期的capacity减1。

现有脚本 #

order_number = None
confirmation_date = None
    
grouped = sales_data['order_number'].unique()
   
# Iterate over the groups and rows within each group
for group in grouped:
    for order_row in range(len(capacity_data)):
        if capacity_data['capacity'][order_row] > 0:
            try:
                order_number.append(group) 
                confirmation_date.append(capacity_data['date'][order_row])
                capacity_data['capacity'][order_row] = capacity_data['capacity'][order_row] - 1
            except:
                 pass
        else:
            pass
orderdict =  dict(zip(order_number, caonfirmation_date))

现有脚本的问题 #

• 初始化错误：order_number和confirmation_date被设为None，调用append方法直接报错，因为None没有该属性。
• 逻辑不符合需求：脚本会从第一个日期开始找可用产能，完全忽略订单本身的order_date，比如订单1003的日期是2022-02-07，脚本会优先占用更早的2022-02-04，不符合“最近可用日期”的要求。
• 效率极低：嵌套循环的时间复杂度为O(N*M)，面对10万+订单时完全无法处理。
• 异常处理滥用：用try-except掩盖错误，导致初始化错误等问题无法被及时发现，后续逻辑直接失效。
• 拼写错误：最后一行的caonfirmation_date是笔误，应为confirmation_date。
• 数据修改不规范：通过capacity_data['capacity'][order_row]直接修改值属于链式索引，会触发SettingWithCopyWarning，且效率低下。

大数据量优化方案 #

核心思路是避免逐行迭代，利用Pandas向量操作和预处理实现高效匹配：

步骤1：预处理产能数据，生成可用日期扩展列表

把每个产能日期按可用次数拆分成单名额行，后续直接按顺序分配：

# 筛选可用产能日期，扩展为单个名额的日期列表
available_dates = []
for _, row in capacity_data[capacity_data['capacity'] > 0].iterrows():
    available_dates.extend([row['date']] * row['capacity'])
# 转换为Series方便后续快速查找
available_dates_series = pd.Series(available_dates)

步骤2：为订单分配符合要求的可用日期

先按订单日期排序，再为每个订单匹配晚于自身日期的最早可用产能：

# 转换订单日期为datetime类型
sales_data['order_date'] = pd.to_datetime(sales_data['order_date'])
# 按订单日期排序，确保先处理早的订单
sales_data_sorted = sales_data.sort_values('order_date').reset_index(drop=True)

confirmation_dates = []
current_idx = 0  # 记录当前可用产能的起始索引
for _, order in sales_data_sorted.iterrows():
    # 从当前索引开始，找第一个晚于订单日期的可用日期
    mask = available_dates_series[current_idx:] > order['order_date']
    if mask.any():
        # 获取第一个符合条件的日期索引
        date_idx = mask.idxmax()
        confirmation_dates.append(available_dates_series[date_idx])
        # 更新起始索引，跳过已占用的产能
        current_idx = date_idx + 1
    else:
        # 无可用产能时的默认处理，这里设为None
        confirmation_dates.append(None)

# 合并结果并还原原订单顺序
sales_data_sorted['confirmation_date'] = confirmation_dates
output = sales_data_sorted.sort_values('order_number')[['order_number', 'confirmation_date']].to_dict('list')
# 将日期转换为字符串格式
output['confirmation_date'] = [d.strftime('%Y-%m-%d') if pd.notna(d) else None for d in output['confirmation_date']]

步骤3：更新产能数据（可选）

如果需要同步更新原capacity_data的剩余产能：

# 统计每个日期的产能被占用次数
occupied_counts = pd.Series(confirmation_dates).value_counts()
# 合并到产能数据表
capacity_data['occupied'] = capacity_data['date'].map(occupied_counts).fillna(0).astype(int)
# 计算剩余产能
capacity_data['remaining_capacity'] = capacity_data['capacity'] - capacity_data['occupied']

优化后的优势 #

• 效率提升显著：预处理为O(M)，分配逻辑为O(N log M)，可轻松处理10万+级别的订单数据。
• 符合需求逻辑：确保分配的是晚于订单日期的最早可用产能。
• 代码规范：所有操作使用Pandas官方API，避免链式索引和错误的异常处理。
• 扩展性强：若后续调整规则（如优先分配某类日期），仅需修改掩码逻辑即可。

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