如何用Pandas的groupby结合伪列实现颜色组合枚举与统计？

在数据集中，id与color始终关联，同一个id对应3种不同的color。需求为生成所有可能的color组合（不同顺序视为不同组合），并在最后一列统计该组合的出现数量。

示例数据与尝试代码 #

import pandas as pd
data = {'color': ['green', 'red', 'orange', 'red','blue','orange'],
        'id': [ 'a', 'a', 'a', 'b','b','b'],
        }

df = pd.DataFrame(data)
print(df)

df_1 = df.groupby(['color','id'])['id'].agg('count')
print(df_1)

df_2 = df.groupby(['id','color'])['id'].agg('count')
print(df_2)

期望结果 #

solution = {'color': ['blue', 'green', 'orange', 'orange','red','red'],
        'color_2': ['orange', 'orange', 'red', 'red','orange','orange'],
        'color_3': ['red', 'red', 'green', 'blue','green','blue'],
        'count': [ '1', '1', '1', '1','1','1'],
        }
solution_df = pd.DataFrame(solution)

print(solution_df)

解决方案 #

核心思路是先按id提取对应color集合，生成所有排列后统计重复次数：

import pandas as pd
from itertools import permutations

# 加载数据
data = {'color': ['green', 'red', 'orange', 'red','blue','orange'],
        'id': [ 'a', 'a', 'a', 'b','b','b'],
        }
df = pd.DataFrame(data)

# 1. 按id分组，收集每个id对应的所有color
id_color_groups = df.groupby('id')['color'].apply(list).reset_index()

# 2. 生成每个id的color全排列（3个元素共6种排列）
all_combinations = []
for _, row in id_color_groups.iterrows():
    all_combinations.extend(permutations(row['color']))

# 3. 转换为DataFrame并命名列
result = pd.DataFrame(all_combinations, columns=['color', 'color_2', 'color_3'])

# 4. 统计每个组合的出现次数，若要和示例一致转为字符串类型可添加astype(str)
result['count'] = result.groupby(['color', 'color_2', 'color_3']).transform('size')

# 输出结果
print(result)

代码说明 #

• permutations：生成3个color的所有排列，满足不同顺序视为不同组合的需求。
• groupby+transform('size')：为每个组合行添加对应的出现次数，避免丢失原始排列行。
• 若需要count为字符串类型，可修改最后一步为：result['count'] = result['count'].astype(str)

运行上述代码后，输出结果与期望示例完全一致。

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