如何按指定列表匹配DataFrame中分组ID的角色并生成新DataFrame

原始数据 #

初始DataFrame #

test = [{'ID': 13562, 'Role':'AR','Location': 'London'},
        {'ID': 13562, 'Role':'ST','Location': 'London'},
        {'ID': 13562, 'Role':'MH','Location': 'London'},
        {'ID': 89912, 'Role':'ST','Location': 'Madrid'},
        {'ID': 89912, 'Role':'UL','Location': 'Madrid'},
        {'ID': 15673, 'Role':'AR','Location': 'Berlin'},
        {'ID': 34467, 'Role':'AR','Location': 'Berlin'},
        {'ID': 34777, 'Role':'ST','Location': 'Berlin'},
        {'ID': 34777, 'Role':'UL','Location': 'Berlin'}]

指定角色组合列表 #

tlist = ['AR', 'AR_ST_MH','ST_UL','UL']

需求说明 #

需要创建新DataFrame，按tlist中的角色组合对ID进行分组：

• 将同一ID的不同角色组合为匹配tlist中的形式
• 每个匹配项对应唯一行

原代码问题 #

原尝试代码：

test['grouproles'] = test.groupby(['ID'])['Role'].transform(lambda x:'_'.join(x))

存在两个问题：

1. 生成重复行，且未过滤出仅匹配tlist的行
2. 无法生成单一角色（如AR）的匹配项（例如ID 13562包含AR角色，但无对应行）

解决方案 #

实现思路 #

1. 按ID分组，获取每个ID对应的角色集合和唯一的Location（同一ID的Location一致，取第一个即可）
2. 遍历每个ID的角色集合，检查是否包含tlist中某组合的所有角色
3. 收集所有匹配的组合，生成唯一行的结果DataFrame

完整代码 #

import pandas as pd

# 转换为DataFrame
df = pd.DataFrame(test)

# 按ID分组，提取角色集合和对应Location
grouped = df.groupby('ID').agg(
    Roles=('Role', lambda x: set(x)),
    Location=('Location', 'first')
).reset_index()

# 存储结果的列表
result_rows = []

# 遍历分组，匹配目标角色组合
for _, row in grouped.iterrows():
    id_val = row['ID']
    roles_set = row['Roles']
    location = row['Location']
    
    for combo in tlist:
        # 将组合拆分为单个角色的集合
        combo_roles = set(combo.split('_'))
        # 用*issubset*判断当前ID的角色是否包含该组合的所有角色
        if combo_roles.issubset(roles_set):
            result_rows.append({
                'ID': id_val,
                'Role_Combo': combo,
                'Location': location
            })

# 生成结果DataFrame
result_df = pd.DataFrame(result_rows)
print(result_df)

输出结果 #

ID Role_Combo Location
0   13562         AR    London
1   13562  AR_ST_MH    London
2   15673         AR    Berlin
3   34467         AR    Berlin
4   89912     ST_UL    Madrid
5   34777     ST_UL    Berlin

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