使用Pandas从分组中系统性选取行生成组合

现有Pandas DataFrame df，简化结构如下（实际包含更多列和group类别）：

group  sub  fruit
a      1    apple
a      2    banana
a      3    orange
b      1    pear
b      2    strawberry
b      3    cherry
c      1    kiwi
c      2    tomato
c      3    lemon

所有group的行数一致。需生成新数据框，满足：

• 每个combo包含全部类型的group
• 同一combo内的sub值无重复

期望输出示例：

combo  group  sub  fruit
1      a      1    apple
1      b      2    strawberry
1      c      3    lemon
2      a      1    apple
2      c      2    tomato
2      b      3    cherry
3      b      1    pear
3      a      2    banana
3      c      3    lemon
4      c      1    kiwi
4      a      2    banana
4      b      3    cherry
5      c      1    kiwi
5      b      2    strawberry
5      a      3    orange
...

错误示例（同一combo内sub重复）：

6      c      2    tomato
6      b      2    strawberry
6      a      3    orange

此前采用随机选取sub的方式，限制过松，无法满足上述要求。

该需求核心是生成group到sub的所有双射组合：每个combo对应一次sub序列的全排列，每个group映射到排列中的唯一sub，同时每个sub仅被一个group使用。

实现代码 #

import pandas as pd
import itertools

# 1. 提取基础信息
groups = df['group'].unique()
subs = df['sub'].unique()

# 2. 生成所有合法的sub排列
sub_permutations = list(itertools.permutations(subs))

# 3. 构建最终数据框
combo_data = []
for combo_id, perm in enumerate(sub_permutations, start=1):
    # 遍历每个group与排列中的对应sub
    for group, sub in zip(groups, perm):
        # 匹配原始数据中的对应行
        matched_row = df[(df['group'] == group) & (df['sub'] == sub)].iloc[0]
        # 组合成包含combo编号的字典
        combo_row = {'combo': combo_id, **matched_row.to_dict()}
        combo_data.append(combo_row)

final_df = pd.DataFrame(combo_data)

# 可选：打乱combo顺序并重新编号
final_df = final_df.sample(frac=1, random_state=42).reset_index(drop=True)
final_df['combo'] = pd.factorize(final_df['combo'])[0] + 1

关键说明 #

• itertools.permutations(subs)生成sub的所有全排列，保证每个排列内sub无重复，从根源上避免同一combo出现重复sub
• 每个排列对应一个combo，通过zip(groups, perm)将每个group与排列中的sub一一绑定，确保每个combo包含所有group
• 直接从原始df中匹配对应行，保留所有原始列的信息，无需额外处理

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