如何按条件分组DataFrame行：将X及后续Y行归为一组

需求：DataFrame的列2包含X、Y、Z三种取值，需要将每个X所在行及其紧随的所有Y行归为一组，同时剔除所有Z行。之前尝试df.groupby(df[2] == 'X')只能单独获取X行，无法实现目标分组。

示例数据 #

import pandas as pd

df = pd.DataFrame({'1':['a','b','c','d','e','f','g','h','i','j','k','l','m','n','o','p'],
                   '2':['Z','X','Y','Z','Z','X','X','Z','X','Y','Y','Z','X','Z','X','Y']})

期望分组结果 #

1  2
1  b  X
2  c  Y
---------
   1  2
5  f  X
---------
   1  2
6  g  X
---------
    1  2
8   i  X
9   j  Y
10  k  Y
---------
    1  2
12  m  X
---------
    1  2
14  o  X
15  p  Y

实现步骤及代码 #

1. 过滤Z行，仅保留X和Y的记录 #

先剔除所有Z行，只留下需要处理的X、Y行：

filtered_df = df[df['2'] != 'Z']

2. 生成分组标识 #

利用cumsum()生成递增的分组ID：每当遇到X时，分组ID加1，后续的Y会沿用当前ID，直到下一个X出现：

group_ids = (filtered_df['2'] == 'X').cumsum()

3. 按分组ID进行分组 #

用生成的group_ids做分组键，就能得到每个X+后续Y的独立分组：

result_groups = filtered_df.groupby(group_ids)

4. 查看分组结果 #

遍历分组即可输出和期望一致的结果：

for _, group in result_groups:
    print(group)
    print('---------')

为什么之前的方法无效？ #

df.groupby(df[2] == 'X')会把所有X行归为一组，**所有非X行（含Y、Z）**归为另一组，完全不符合“每个X+紧随Y独立成组”的需求，必须用cumsum()生成动态递增的分组ID来实现目标。

内容的提问来源于stack exchange，提问作者doine