将 pandas groupby() 中的值提取到结合单个值和 numpy 数组的新数据集中

Question

我有一个名为df的熊猫数据帧，看起来像这样

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分组的数据集，并按test_type提取correct的平均值(作为单个值)以及按test_type和test_number正确(作为numpy数组)的平均值。

这是我需要的：

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

我一直在使用df.groupby(["name", "test_type"]).correct.mean().reset_index()和df.groupby(["name", "test_type", "test_number"]).correct.mean().reset_index()但我无法设法 1) 像我想要的那样通过test_number提取平均值作为数组和 2) 在连贯的数据帧中组织输出。

提前谢谢。

Answer 1

IIUC，您可以使用：

A = df.groupby(['name', 'test_type'], sort=False)['correct'].mean().unstack()
B = (df
.groupby(['name', 'test_type', 'test_number'])['correct'].mean()
.unstack().agg(list, axis=1).unstack()
)
out = A.join(B.add_suffix('_by_tn')).add_prefix('correct_')

输出：

test_type  correct_0  correct_1 correct_0_by_tn correct_1_by_tn
name                                                           
joe             0.75       0.50      [1.0, 0.5]      [0.0, 1.0]
jim             0.50       0.25      [0.0, 1.0]      [0.0, 0.5]

替代输出：

out = (A
.join(B.add_suffix('_by_tn'))
.add_prefix('correct_')
.rename_axis(columns=None)
.reset_index()
)

输出：

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