用pandas高效实现按库存标记补货通知字段的最优方案

实现方案 #

你提到的用累加计数的思路是完全可行的，全程用pandas向量化操作实现，比嵌套循环效率高很多，代码如下：

前置修正

你给出的库存表示例中PRODUCT_ID存在笔误，写为9977和用户表的9777不匹配，下方代码先修正该问题保证匹配逻辑正确。

完整代码

import pandas as pd
import numpy as np

# 原用户登记表
df = pd.DataFrame([['nick_1@gmail.com', 9999, 1000, 1001, 1002, np.nan ], 
                   ['sam_s@gmail.com', 9999, 1000, 1001, 1002, np.nan ],
                   ['george_l@gmail.com', 9999, 1000, 1001, 1500, np.nan ],
                   ['fred_g@gmail.com', 9777, 1500, 1501, 1155, np.nan ],
                   ['katie_l@gmail.com', 9777, 1500, 1501, 1155, np.nan ]], 
columns=['CUSTOMER_EMAIL', 'PRODUCT_ID', 'STORE_1', 'STORE_2', 'STORE_3', 'SEND_NOTIFICATION'])

# 修正笔误后的库存表
stock = pd.DataFrame([[9999, 1000, 2], 
                      [9999, 1002, 1 ],
                      [9777, 1500, 1 ]
                   ], columns=['PRODUCT_ID', 'STORE_ID', 'STOCK_AVAILABLE'])

# 步骤1：保留原始用户登记顺序索引
df['orig_idx'] = df.index

# 步骤2：宽表转长表，把3个门店列展开为每行对应一个用户+商品+单个门店
user_store_long = df.melt(
    id_vars=['orig_idx', 'CUSTOMER_EMAIL', 'PRODUCT_ID'],
    value_vars=['STORE_1', 'STORE_2', 'STORE_3'],
    value_name='STORE_ID'
).drop('variable', axis=1)

# 步骤3：关联库存表，获取每个用户对应门店的可用库存
user_store_stock = user_store_long.merge(
    stock,
    on=['PRODUCT_ID', 'STORE_ID'],
    how='left'
)

# 步骤4：过滤有库存的记录，按登记顺序排序后分组累计计数，筛选出库存在配额内的用户
valid_records = user_store_stock[user_store_stock['STOCK_AVAILABLE'].notna()].sort_values('orig_idx')
valid_records['group_cum_count'] = valid_records.groupby(['PRODUCT_ID', 'STORE_ID']).cumcount() + 1
valid_user_idx = valid_records[valid_records['group_cum_count'] <= valid_records['STOCK_AVAILABLE']]['orig_idx'].unique()

# 步骤5：回写结果到原表
df['SEND_NOTIFICATION'] = df.index.isin(valid_user_idx)
# 删除辅助列
df = df.drop('orig_idx', axis=1)

print(df)

输出结果

完全符合你给出的规则要求。

逻辑说明

groupby + cumcount本质就是你猜测的cumsum实现思路，等价于分组内逐行累加1，正好可以用来统计每个商品+门店组合下匹配到的用户顺序，只要序号小于等于可用库存数，就属于可通知用户，最后对符合条件的用户索引去重避免重复标记即可。

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