基于Race_ID分组生成Student_ID两两组合新DataFrame的需求

现有一个以Race_ID为索引的DataFrame，数据如下：

Race_ID        Student_ID      Rank  
1              1               3     
1              2               2     
1              3               1     
1              4               4     
2              1               2     
2              2               3     
2              3               1     

需基于该DataFrame创建新的DataFrame，规则为：在每个Race_ID分组内，将Student_ID进行两两组合生成新行，并计算对应学生的Rank之和（Rank_sum）。预期输出的DataFrame如下：

Race_ID   Student_ID_pair   Rank_sum
1         1-2               5
1         1-3               4        
1         1-4               7
1         2-3               3
1         2-4               6
1         3-4               5
2         1-2               5
2         1-3               3
2         2-3               4

其中Race_ID=1对应4选2的6行，Race_ID=2对应3选2的3行。

可以用pandas结合itertools.combinations实现需求，代码如下：

import pandas as pd
from itertools import combinations

# 构造原始DataFrame
data = {
    'Race_ID': [1,1,1,1,2,2,2],
    'Student_ID': [1,2,3,4,1,2,3],
    'Rank': [3,2,1,4,2,3,1]
}
df = pd.DataFrame(data).set_index('Race_ID')

# 重置索引，将Race_ID转为列
df_reset = df.reset_index()

# 定义分组处理函数
def process_group(group):
    student_rank_pairs = list(zip(group['Student_ID'], group['Rank']))
    # 生成所有不重复的两两组合
    for pair in combinations(student_rank_pairs, 2):
        (s1, r1), (s2, r2) = pair
        yield {
            'Race_ID': group['Race_ID'].iloc[0],
            'Student_ID_pair': f"{s1}-{s2}",
            'Rank_sum': r1 + r2
        }

# 分组处理并拼接结果
result_rows = []
for _, group in df_reset.groupby('Race_ID'):
    result_rows.extend(process_group(group))

final_df = pd.DataFrame(result_rows)
print(final_df)

关键步骤说明 #

1. 重置索引：将原DataFrame的Race_ID索引转为列，便于后续分组操作；
2. 分组生成组合：对每个Race_ID分组，用combinations生成所有学生的不重复两两组合；
3. 计算与拼接：遍历组合，生成Student_ID_pair字符串并计算Rank_sum，最后整理成目标DataFrame。

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