如何将pandas groupby()结果提取为含单值与numpy数组的新数据集

原始数据 #

我们有一个名为df的pandas DataFrame，数据如下：

name   test_type   test_number   correct
joe    0           1             1
joe    0           2             0
joe    1           1             0
joe    1           2             1
joe    0           1             1
joe    0           2             1
jim    1           1             0
jim    1           2             1
jim    0           1             0
jim    0           2             1
jim    1           1             0
jim    1           2             0

需求 #

按name分组，提取两类统计结果并整合成目标DataFrame：

1. 按test_type分组的correct均值（单值形式，列名如correct_0、correct_1）
2. 按test_type和test_number分组的correct均值（numpy数组形式，列名如correct_0_by_tn、correct_1_by_tn）

目标输出：

name    correct_0    correct_1    correct_0_by_tn    correct_1_by_tn
joe     0.75         0.5          [1, 0.5]           [0, 1]
jim     0.5          0.25         [0, 1]             [0, 0.5]

解决方案 #

通过两次分组聚合+结果合并的方式实现，步骤如下：

1. 计算test_type维度的均值并转宽表 #

先统计每个name+test_type的均值，再用unstack将行转列，得到单值形式的统计列：

# 计算均值并转宽表
mean_by_type = df.groupby(["name", "test_type"])["correct"].mean().unstack()
# 重命名列以匹配需求
mean_by_type.columns = [f"correct_{col}" for col in mean_by_type.columns]
mean_by_type = mean_by_type.reset_index()

2. 计算test_type+test_number维度的均值并转为数组 #

先统计name+test_type+test_number的均值，再按name+test_type分组，将每组的均值转为numpy数组，同样转成宽表：

# 计算三级分组的均值
mean_by_type_tn = df.groupby(["name", "test_type", "test_number"])["correct"].mean()
# 将每组均值转为numpy数组并转宽表
mean_by_type_tn = mean_by_type_tn.groupby(["name", "test_type"]).apply(lambda x: x.to_numpy()).unstack()
# 重命名列
mean_by_type_tn.columns = [f"correct_{col}_by_tn" for col in mean_by_type_tn.columns]
mean_by_type_tn = mean_by_type_tn.reset_index()

3. 合并结果表 #

将两个统计结果按name字段合并，得到最终结构：

final_df = mean_by_type.merge(mean_by_type_tn, on="name")

完整可运行代码 #

import pandas as pd
import numpy as np

# 构造原始DataFrame
data = [
    ["joe", 0, 1, 1],
    ["joe", 0, 2, 0],
    ["joe", 1, 1, 0],
    ["joe", 1, 2, 1],
    ["joe", 0, 1, 1],
    ["joe", 0, 2, 1],
    ["jim", 1, 1, 0],
    ["jim", 1, 2, 1],
    ["jim", 0, 1, 0],
    ["jim", 0, 2, 1],
    ["jim", 1, 1, 0],
    ["jim", 1, 2, 0]
]
df = pd.DataFrame(data, columns=["name", "test_type", "test_number", "correct"])

# 处理test_type维度的均值
mean_by_type = df.groupby(["name", "test_type"])["correct"].mean().unstack()
mean_by_type.columns = [f"correct_{col}" for col in mean_by_type.columns]
mean_by_type = mean_by_type.reset_index()

# 处理test_type+test_number维度的数组均值
mean_by_type_tn = df.groupby(["name", "test_type", "test_number"])["correct"].mean()
mean_by_type_tn = mean_by_type_tn.groupby(["name", "test_type"]).apply(lambda x: x.to_numpy()).unstack()
mean_by_type_tn.columns = [f"correct_{col}_by_tn" for col in mean_by_type_tn.columns]
mean_by_type_tn = mean_by_type_tn.reset_index()

# 合并得到最终结果
final_df = mean_by_type.merge(mean_by_type_tn, on="name")
print(final_df)

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