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

时间:2016-04-06 18:14:18

标签: python pandas

我有一个按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')。head(3).sum()

但这并不是我需要的。任何帮助,建议,代码片段都会非常感激。

2 个答案:

答案 0 :(得分:3)

  

我试图玩

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])

有关split-apply-combine方法的更多信息,请参阅here

答案 1 :(得分:2)

您可以先设置grp_idx

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

现在idgrp_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排序为正确。