python - Pandas :通过聚合折叠每组中的前 n 行

Question

我有一个按 id 分组的数据框。有很多组，每个组都有可变数量的行。所有组的前三行不包含有趣的数据。我想按以下方式“折叠”每组中的前三行以形成一行:

“id”和“type”在新的“collapsed”行中将保持不变。
‘grp_idx’在前三行发生聚合时将重命名为“0”
col_1 将是前三行的总和
col_2 将是前三行的总和
如果前 3 行中的值全部为 0，则“折叠”行中的“标志”将为 0。如果前三行中的任何一行为 1，则 'flag' 将为 1。 (一个简单的总和就足够了这个逻辑，因为标志只在一行中为所有组设置)

这是数据框的示例:

import pandas as pd
import numpy as np   
df = pd.DataFrame.from_items([
    ('id', [283,283,283,283,283,283,283,756,756,756]), 
    ('type', ['A','A','A','A','A','A','A','X','X','X']),
    ('grp_idx', [1,2,3,4,5,6,7,1,2,3]),
    ('col_1', [2,4,6,8,10,12,14,5,10,15]),
    ('col_2', [3,6,9,12,15,18,21,1,2,3]),
    ('flag', [0,0,0,0,0,0,1,0,0,1]),
    ]);
print(df)

    id   type  grp_idx  col_1  col_2  flag
0  283    A        1      2      3     0
1  283    A        2      4      6     0
2  283    A        3      6      9     0
3  283    A        4      8     12     0
4  283    A        5     10     15     0
5  283    A        6     12     18     0
6  283    A        7     14     21     1
7  756    X        1      5      1     0
8  756    X        2     10      2     0
9  756    X        3     15      3     1

处理后，我希望数据框看起来像:

ID  Type   grp_idx  col_1  col_2   flag
283  A         0     12      18      0
283  A         4     8       12      0
283  A         5     10      15      0
283  A         6     12      18      0
283  A         7     14      21      1
756  X         0     30       6      1

我不确定如何进行。我试着玩弄

df.groupby('id').head(3).sum()

但这不是我需要的。非常感谢任何帮助、建议和代码片段。

Answer 1

I was trying to play around with

df.groupby('id').head(3).sum()

在调用groupby() 之后，您需要aggregate() 以便按照您想要的方式组合。尝试这样的事情:

# function to sum the first 3 rows
def head_sum(x):
    return x.head(3).sum()

# function to get max of first 3 rows
def head_max(x):
    return x.head(3).max()

# We can use a dictionary in `aggregate()` to call a 
# specific function for each column in the groupby
column_funcs = {'col_1': head_sum,
                'col_2': head_sum,
                'flag': head_max,
                'id': max,  # all the vals should be the same
                'type': max}  # are the 'id' and 'type' always matched?
collapsed = df.groupby('id').aggregate(column_funcs)
collapsed['grp_idx'] = 0

new_df = pd.concat([df, collapsed])

参见 here有关拆分-应用-组合方法的更多信息。