如何按Team_members无序列表对Pandas DataFrame进行分组聚合

完整可运行代码 #

import pandas as pd
import numpy as np

# 测试数据集
data = {'Day': [1, 1, 1, 1, 2, 2, 2, 2],
        'Worker_ID': [1, 2, 3, 4, 1, 2, 3, 4],
        'Skills': ['[1 3]', '[2 5]', '[4 2]', '[3 3]', '[2 4]', '[3 5]', '[4 3]', '[2 2]'],
        'Team_members': ['[1 3]', '[4 2]', '[3 1]', '[2 4]', '[1 3]', '[4 2]', '[3 1]', '[2 4]']}
df = pd.DataFrame(data)

# 1. 生成原始数据副本避免修改原数据
df_copy = df.copy()

# 定义辅助函数：将"[x y]"格式的字符串转换为整数列表
def str_to_list(input_str):
    return list(map(int, input_str.strip('[]').split()))

# 2. 标准化团队成员列：排序成员ID消除顺序差异，相同成员的不同排列会被识别为同一个团队
df_copy['Team_members'] = df_copy['Team_members'].apply(str_to_list)
df_copy['Team_members'] = df_copy['Team_members'].apply(sorted)
# 转回字符串格式方便后续分组与输出
df_copy['Team_members'] = df_copy['Team_members'].apply(lambda x: f"[{x[0]} {x[1]}]")

# 转换技能列为numpy数组，方便后续逐位计算均值
df_copy['Skills'] = df_copy['Skills'].apply(str_to_list).apply(np.array)

# 3. 按日期+团队成员分组，计算团队平均技能
grouped_df = df_copy.groupby(['Day', 'Team_members']).agg(
    Team_Skills=('Skills', lambda skill_arr: np.mean(np.vstack(skill_arr), axis=0))
).reset_index()

# 4. 为唯一的团队组合分配自增ID
team_id_map = {team: idx+1 for idx, team in enumerate(grouped_df['Team_members'].unique())}
grouped_df['Team_ID'] = grouped_df['Team_members'].map(team_id_map)

# 5. 格式化技能列为要求的字符串格式，调整列顺序得到最终结果
grouped_df['Team_Skills'] = grouped_df['Team_Skills'].apply(lambda skill: f"[{skill[0]} {skill[1]}]")
final_result = grouped_df[['Day', 'Team_ID', 'Team_Skills', 'Team_members']]

print(final_result)

输出结果匹配说明 #

运行上述代码得到的输出与要求的格式完全一致：

Day  Team_ID Team_Skills Team_members
0    1        1  [2.5 2.5]        [1 3]
1    1        2    [2.5 4.0]        [2 4]
2    2        1    [3.0 3.5]        [1 3]
3    2        2  [2.5 3.5]        [2 4]

如果真实数据中团队成员数量、技能维度有变化，仅需调整辅助函数的格式化逻辑，核心分组、计算逻辑无需修改，适配大型数据集处理。

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