如何按多条件及分组为DataFrame设置产品映射值？

问题背景 #

现有按ID分组的DataFrame，每组固定包含7条产品记录（存在重复产品可能），需按以下规则完成产品到PRODUCT_1至PRODUCT_5的映射，剩余2个产品映射值为空：

• 规则1：RANK为1、2、3的产品，分别映射为PRODUCT_1、PRODUCT_2、PRODUCT_3
• 规则2：若YELLOW不在排名前三，为其随机分配PRODUCT_4或PRODUCT_5
• 规则3：从未被映射的剩余产品中随机选1个，分配剩下的PRODUCT_4或PRODUCT_5
• 规则4：最后2个未被选中的产品映射值为空

原始数据 #

预期结果示例 #

解决方案（Python Pandas实现） #

import pandas as pd
import numpy as np

# 加载原始数据
df = pd.DataFrame({
    'ID': ['A']*7 + ['B']*7,
    'PRODUCT': ['RED', 'BLUE', 'BLACK', 'GREEN', 'GREEN', 'PURPLE', 'YELLOW',
                'PURPLE', 'BLACK', 'BLACK', 'GREEN', 'WHITE', 'BROWN', 'GREEN'],
    'RANK': [1,2,3,4,5,6,7,1,2,3,4,5,6,7]
})

# 初始化MAP列为空字符串
df['MAP'] = ""

def map_products(group):
    # 规则1：处理排名前三的产品
    group.loc[group['RANK'] == 1, 'MAP'] = 'PRODUCT_1'
    group.loc[group['RANK'] == 2, 'MAP'] = 'PRODUCT_2'
    group.loc[group['RANK'] == 3, 'MAP'] = 'PRODUCT_3'
    
    # 获取已映射的索引
    mapped_idx = group[group['RANK'].isin([1,2,3])].index
    # 筛选未映射的记录
    remaining_idx = group[~group.index.isin(mapped_idx)].index
    yellow_in_top3 = (group.loc[mapped_idx, 'PRODUCT'] == 'YELLOW').any()
    
    assigned_tags = []
    # 规则2：处理不在前三的YELLOW
    if not yellow_in_top3:
        yellow_rows = group.loc[remaining_idx][group['PRODUCT'] == 'YELLOW']
        if not yellow_rows.empty:
            yellow_tag = np.random.choice(['PRODUCT_4', 'PRODUCT_5'])
            group.loc[yellow_rows.index, 'MAP'] = yellow_tag
            assigned_tags.append(yellow_tag)
            remaining_idx = remaining_idx.difference(yellow_rows.index)
    
    # 规则3：分配剩余的PRODUCT_4/5标签
    if len(assigned_tags) == 0:
        # YELLOW在前三时，直接随机选两个产品分配两个标签
        selected_idx = np.random.choice(remaining_idx, 2, replace=False)
        group.loc[selected_idx[0], 'MAP'] = 'PRODUCT_4'
        group.loc[selected_idx[1], 'MAP'] = 'PRODUCT_5'
    else:
        # 取未使用的标签，随机选一个剩余产品分配
        remaining_tag = 'PRODUCT_4' if assigned_tags[0] == 'PRODUCT_5' else 'PRODUCT_5'
        if len(remaining_idx) > 0:
            selected_idx = np.random.choice(remaining_idx)
            group.loc[selected_idx, 'MAP'] = remaining_tag
    
    return group

# 按ID分组执行映射逻辑
result_df = df.groupby('ID', group_keys=False).apply(map_products)

# 输出结果
print(result_df)

代码逻辑说明 #

1. 数据初始化：加载原始数据并将MAP列初始化为空字符串
2. 分组处理：对每个ID组单独执行映射规则：
   • 先为排名1-3的产品直接分配对应标签
   • 判断YELLOW是否在前三，若不在则随机分配PRODUCT_4或PRODUCT_5
   • 从未被映射的剩余产品中随机挑选一个，分配剩下的PRODUCT_4或PRODUCT_5
   • 剩余未被选中的产品保持MAP为空（已初始化，无需额外操作）
3. 随机逻辑：使用np.random.choice实现随机分配，保证每次运行结果可能存在差异（符合规则要求）

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