Pandas：按聚合方式折叠每个组中的前n行

Question

我有一个数据框，按id分组。有很多组，每个组的行数不确定。所有组的前三行都不包含有用的数据。我想要将每个组的前三行“合并”成一行，如下所示：

在新的“合并”行中，“id”和“type”保持不变。 当聚合前三行时，“grp_idx”将重命名为“0”。 col_1将是前三行的总和。 col_2将是前三行的总和。 如果前三行的值都为0，则“collapsed”行中的“flag”为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()

但这并不是我需要的。如果有帮助、建议或代码片段，将不胜感激。

- Learner

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

关于分组-应用-合并方法的更多信息，请查看此处。

- Zachary Cross

---

谢谢Zachary！你提出的技巧真的很好。这是我将添加到我薄弱知识中的东西。你的建议肯定可以完成所需的工作。我只需要删除前三行。再次感谢你分享你的知识。 - Learner

Answer 2

您可以通过设置 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对其进行排序以使其正确。

- ayhan

---

太棒了！我自己永远也想不出来。感谢你抽出时间来帮助我。 - Learner