在熊猫中使用 2 列的自定义聚合函数

时间:2021-04-02 15:37:37

标签: python pandas dataframe

我有下表

event_name | score | date      | flag | 
event_1    | 123   | 12APR2018 |  0   |
event_1    | 34    | 05JUN2019 |  0   |
event_1    | 198   | 08APR2020 |  0   |
event_2    | 3     | 14SEP2019 |  0   |
event_2    | 34    | 22DEC2019 |  1   |
event_2    | 90    | 17FEB2020 |  0   | 
event_3    | 772   | 19MAR2021 |  1   |

我想获得

event_name | sum_score | date_flag_1 | 
event_1    | 355       |             | 
event_2    | 127       | 22DEC2019   | 
event_3    | 772       | 19MAR2021   | 

其中 sum_score 是对应事件的列分数总和,date_flag_1 是对应事件的 flag = 1 时的第一个日期。如果当前事件的所有行的 flag = 0,date_flag_1 应该缺失

我想代码应该看起来像

df_agg = df.groupby('event_name').agg({'score': 'sum', ['date', 'flag']: my_custom_function})
df_agg.columns = ['event_name', 'sum_score', 'date_flag_1']

但是,我不确定应该如何实现 my_custom_function,这将是一个自定义聚合函数,它使用两列而不是一列(与其他聚合函数一样)。请帮忙

1 个答案:

答案 0 :(得分:2)

聚合两次,concat 结果。您可以设置子集的第二个聚合然后使用内置的 GroupBy.first

import pandas as pd

pd.concat([df.groupby('event_name')['score'].sum(),
           df[df.flag.eq(1)].groupby('event_name')['date'].first().rename('date_flag_1')], 
          axis=1)

#            score date_flag_1
#event_name                   
#event_1       355         NaN
#event_2       127   22DEC2019
#event_3       772   19MAR2021

举例来说,这可以通过一个 agg 调用来完成;但是它会很慢,因为这需要一个 lambda x:,它将被计算为组上的慢循环(与矢量化/cythonized 内置 GroupBy 操作相反)。

因为 .agg 仅作用于单个系列,所以hacky 的解决方法是创建一个同时接受系列和数据帧的函数。您使用 Series 索引对 DataFrame 进行子集化(您必须有一个非重复索引才能正常工作),然后您可以进行可以使用多列的聚合。这既过于复杂又缓慢,所以我不会这样做。

def get_first_date(s, df):
    # rows within group where `s==1`
    res = df.loc[s[s.eq(1)].index, 'date'].dropna()

    if not res.empty:
        return res.iloc[0]
    else:
        return np.NaN

df.groupby('event_name').agg({'score': 'sum', 
                              'flag': lambda x: get_first_date(x, df)})

#            score       flag
#event_name                  
#event_1       355        NaN
#event_2       127  22DEC2019
#event_3       772  19MAR2021