python - Pandas groupby 并应用 - 通过 groupby 变量获取新的 DataFrame

Question

我正在尝试使用 pandas.DataFrame.groupby['x'] 对分组的 df 进行计算，按 x.

'x' 重复多次时会出现问题。 apply 函数将执行 'x' 重复的计算次数，尽管我只需要“聚合”值(它不是真正的聚合，但更像是 - 处理)。

这是一个玩具示例:

def simulate_complicated_func(df):
    # This function simulates complicate calculations
    returned_col_names = ['calc1', 'calc2', 'calc3']

    df['calc1'] = ''.join(df['var1'])
    df['calc2'] = df['var2'].mean()
    df['calc3'] = ''.join(df['var1']) + str(df['var2'].max())

    return df[['id'] + returned_col_names]

df = pd.DataFrame({'id':['id1', 'id1', 'id2', 'id3', 'id3', 'id3'],
                   'var1':['abc', 'cba', 'abc', 'cba', 'abc', 'cba'],
                   'var2':[9, 4, 7, 4, 1, 3]})

print(df)

    id var1  var2
0  id1  abc     9
1  id1  cba     4
2  id2  abc     7
3  id3  cba     4
4  id3  abc     1
5  id3  cba     3

res_df = df.groupby(['id']).apply(simulate_complicated_func).drop_duplicates()
print(res_df)

    id      calc1     calc2       calc3
0  id1     abccba  6.500000     abccba9
2  id2        abc  7.000000        abc7
3  id3  cbaabccba  2.666667  cbaabccba4

输出正是我想要的，但效率不高。使用 pandas 有更好的方法吗？

编辑:如何优化？

如果我们将 print 语句添加到 simulate_complicated_func()

def simulate_complicated_func(df):
    # This function simulates complicate calculations
    print("function called")
    # ...

我们可以看到代码会打印 6 次:

function called
function called
function called
function called
function called
function called

实际上，我们只需要访问这个函数3次(groupby创建的组数)。

Answer 1

一个想法是从自定义函数返回 Series，因此 drop_duplicates 不是必需的:

def simulate_complicated_func(df):
    # This function simulates complicate calculations
    returned_col_names = ['calc1', 'calc2', 'calc3']

    a = ''.join(df['var1'])
    b = df['var2'].mean()
    c = ''.join(df['var1']) + str(df['var2'].max())

    return pd.Series([a,b,c], index=returned_col_names)

res_df = df.groupby(['id']).apply(simulate_complicated_func).reset_index()
print(res_df)
    id      calc1     calc2       calc3
0  id1     abccba  6.500000     abccba9
1  id2        abc  7.000000        abc7
2  id3  cbaabccba  2.666667  cbaabccba4

另一个想法是使用 DataFrameGroupBy.agg但它只能用于处理具有聚合函数(如 join 和 mean)的所有列。函数 agg 分别处理每一列，因此 cal3 不是简单/有效的计数方式 - 再次需要自定义函数和最后一个连接输出:

def simulate_complicated_func(df):
    # This function simulates complicate calculations
    returned_col_names = ['calc3']
    c = ''.join(df['var1']) + str(df['var2'].max())
    return pd.Series([c], index=returned_col_names)

d = {'var1': ''.join, 'var2':'mean'}
cols = {'var1':'calc1','var2':'calc2'}
g = df.groupby(['id'])

df1 = g.agg(d).rename(columns=cols)
print (df1)
         calc1     calc2
id                      
id1     abccba  6.500000
id2        abc  7.000000
id3  cbaabccba  2.666667

df2 = df.groupby(['id']).apply(simulate_complicated_func)
print(df2)
          calc3
id             
id1     abccba9
id2        abc7
id3  cbaabccba4

df = pd.concat([df1, df2], axis=1).reset_index()
print (df)
    id      calc1     calc2       calc3
0  id1     abccba  6.500000     abccba9
1  id2        abc  7.000000        abc7
2  id3  cbaabccba  2.666667  cbaabccba4