如何在创建多个新列的pandas.core.groupby.GroupBy函数上运行单个聚合函数

时间:2018-10-29 12:51:56

标签: python pandas

请考虑以下数据框:

df = pd.DataFrame({
    'group': [i % 3 for i in range(10)],
    'a': np.random.rand(10),
    'b': np.random.rand(10)
})

def my_agg(x):
    x = x.values.reshape([x.shape[0] // 2,2])
    prod = x[:,0] * x[:,1]
    return [np.sum(prod), np.mean(prod)]

df.set_index('group').stack().groupby('group').apply(my_agg)

作为结果生产

group
0     [0.3625660911145343, 0.09064152277863358]
1       [1.132618561193485, 0.3775395203978283]
2    [0.37300784663400804, 0.12433594887800269]
dtype: object

而我希望每一列都有单独的列。考虑到以下几点,在大熊猫中有没有一种整齐的方法可以做到这一点?

  • 生成的多个特征更加复杂,并且一起计算它们的效率更高;
  • 功能数量远大于2?

1 个答案:

答案 0 :(得分:1)

您可以将输出转换为list,然后由构造函数转换为DataFrame

def my_agg(x):
    x = x.values.reshape([x.shape[0] // 2,2])
    return [np.sum(x[:,0] * x[:,1]), np.mean(x[:,0] * x[:,1])]

s = df.set_index('group').stack().groupby('group').apply(my_agg)
df1 = pd.DataFrame(s.values.tolist(), index=s.index, columns=['a','b'])
print (df1)
              a         b
group                    
0      2.210601  0.552650
1      0.335913  0.111971
2      1.696796  0.565599

或者您可以先返回Series然后再返回unstack,但这应该更慢:

def my_agg(x):
        x = x.values.reshape([x.shape[0] // 2,2])
        return pd.Series([np.sum(x[:,0] * x[:,1]), np.mean(x[:,0] * x[:,1])], index=['a','b'])

df1 = df.set_index('group').stack().groupby('group').apply(my_agg).unstack()
print (df1)
              a         b
group                    
0      0.391921  0.097980
1      0.417366  0.139122
2      0.788845  0.262948