如何在Python中计算DataFrame分组各维度的平均值

给定如下DataFrame，其中vector列每行是一个多维数组，需求是按group分组，计算每组内所有向量各维度的平均值：

import pandas as pd

df = pd.DataFrame({
    'user': ['user_1', 'user_2', 'user_3', 'user_4', 'user_5', 'user_6'],
    'vector': [[1, 0, 2, 0], [1, 8, 0, 2], [6, 2, 0, 0], [5, 0, 2, 2], [3, 8, 0, 0], [6, 0, 0, 2]],
    'group': ['A', 'B', 'C', 'B', 'A', 'A']
})

解决方案 #

方法一：拆分向量列后计算均值 #

将vector列拆分为独立数值列，分组计算均值后再合并为列表：

# 拆分vector为多列
vector_cols = pd.DataFrame(df['vector'].tolist(), index=df.index)
# 合并分组列与拆分后的向量列
combined = pd.concat([df['group'], vector_cols], axis=1)
# 按group分组计算各维度均值，保留1位小数
group_mean = combined.groupby('group').mean().round(1)
# 将每行转换为列表格式
result = group_mean.apply(list, axis=1)

# 输出结果
for group, avg_vec in result.items():
    print(f"group {group}: {avg_vec}")

方法二：直接对分组后的向量列表计算均值 #

借助numpy直接处理分组后的向量列表，实现更简洁：

import numpy as np

# 按group分组，计算每组向量的维度均值并保留1位小数
grouped_avg = df.groupby('group')['vector'].apply(
    lambda x: np.mean(x.tolist(), axis=0).round(1).tolist()
)

# 输出结果
for group, avg_vec in grouped_avg.items():
    print(f"group {group}: {avg_vec}")

输出结果 #

group A: [3.3, 2.7, 0.7, 0.7]
group B: [3.0, 4.0, 1.0, 2.0]
group C: [6.0, 2.0, 0.0, 0.0]

注：若需匹配示例中保留一位小数的近似值（如2.6而非2.7），可调整round()参数或手动截取小数位。

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