Python pandas用自定义条件分组、转换多个列

Question

我有一个包含500k+记录的数据框，我想按多列(字符串和日期的数据类型)分组，然后根据自定义条件在每个组中只选择少数记录。

基本上，我需要对记录进行分组(按first_roll_up,date,granular_timestamp)以检查该组是否包含列top的任何值，如果存在，则仅选择top值的记录。同样，如果组中没有top值的记录，则选择所有记录。

输入:

<表类>first_roll_up子顶部日期granular_timestamp值tbody><<tr>ABCT12/10/20222/10/2022 10:00:00:000。ABCSUB_A_12/10/20222/10/2022 10:00:00:000。ABCSUB_A_22/10/20222/10/2022 10:00:00:000。ABCSUB_A_32/10/20222/10/2022 10:00:00:000。XYZSUB_X_12/12/20222/10/2022 11:00:00:000。XYZSUB_X_22/12/20222/10/2022 11:00:00:000。XYZSUB_Y_12/12/20222/10/2022 12:00:00:000。

这里似乎根本不需要使用groupby()，而是需要一些过滤和连接。

首先，我们将创建一个数据框，它可以帮助我们识别一行的top值是否为T1。

df_t1 = df[df['top'] == 'T1']

我们现在想要找到所有而不是包含T1值的行，并通过first_roll_up将其连接起来。这似乎只发生在sub缺失时，这使得您的数据看起来是分层的(例如first_roll_up—>sub)。

df_want = pd.merge(df, df_t1, on='first_roll_up', how='outer', indicator=True, suffixes=('', '_drop')) 
.query("_merge == 'left_only'") 
.drop(columns='_merge')

这给了我们一个像这样的数据框架:

first_roll_up      sub  top  ... date_drop granular_timestamp_drop values_drop
4           XYZ  SUB_X_1       ...       NaN                     NaN         NaN
5           XYZ  SUB_X_2       ...       NaN                     NaN         NaN
6           XYZ  SUB_Y_1       ...       NaN                     NaN         NaN

由于我们连接的方式，我们有一组我们不关心的列。我们用_drop作为它们的后缀。我们稍后会删除它们。

现在我们将T1的值堆栈在我们刚刚创建的数据框上并删除我们不想要的列:

df_want = pd.concat([df_t1, df_want.loc[:, ~df_want.columns.str.endswith('_drop')]], ignore_index=True)

最终输出:

first_roll_up      sub top       date      granular_timestamp values
0           ABC           T1  2/10/2022  2/10/2022 10:00:00:000      .
1           XYZ  SUB_X_1      2/12/2022  2/10/2022 11:00:00:000      .
2           XYZ  SUB_X_2      2/12/2022  2/10/2022 11:00:00:000      .
3           XYZ  SUB_Y_1      2/12/2022  2/10/2022 12:00:00:000      .

另一个可能的解决方案:

(df.groupby(['first_roll_up', 'date', 'granular_timestamp'], as_index=False)
.apply(lambda g: g if g.top.isna().all() else g.loc[g.top.notna()])
.reset_index(drop=True))

输出:

first_roll_up       sub  top        date       granular_timestamp values
0          ABC        NaN  T1   2/10/2022   2/10/2022 10:00:00:000       .
1          XYZ   SUB_X_1   NaN  2/12/2022   2/10/2022 11:00:00:000       .
2          XYZ   SUB_X_2   NaN  2/12/2022   2/10/2022 11:00:00:000       .
3          XYZ   SUB_Y_1   NaN  2/12/2022   2/10/2022 12:00:00:000       .