python – Pandas：通过聚合折叠每组中的前n行

我有一个按ID分组的数据框.有许多组,每组具有可变数量的行.所有组的前三行不包含有趣的数据.我想按照以下方式“折叠”每组中的前三行以形成一行：

‘id’和’type’在新的’折叠’行中保持不变.
当前三行的聚合发生时,’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’).头部(3)的.sum()

但这不是我需要的.任何帮助,建议,代码片段将非常感激.

解决方法:

您可以从设置grp_idx开始：

df["grp_idx"] = np.where(df.groupby("id").cumcount()<3, 0, df["grp_idx"])

现在id和grp_idx创建你想要的分组：

df.groupby(["id", "type", "grp_idx"]).sum().reset_index()

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

我假设类型不能与同一个id不同,因为您没有为该列提供任何条件.我还假设df按id排序.如果没有,您可以先为grp_idx排序,使其正确.