如何按Pandas DataFrame中b列的x起始z终止规则进行分组？

原始DataFrame #

import pandas as pd

df = pd.DataFrame(
    {
        'a': [1, 10, 20, 30, 40, 90, 100, 200, 11], 
        'b': ['x', 'y', 'h', 'z', 'x', 'z', 'x', 'a', 'z']
    }
)

分组需求 #

需要按以下规则分组：

• 当b列出现x时开始分组
• 当b列出现z时结束分组
• 分组需包含x到z之间的所有行，示例分组结果如下：

0    1  x
1   10  y
2   20  h
3   30  z

4   40  x
5   90  z

6  100  x
7  200  a
8   11  z

解决方案 #

可以通过标记分组起止位置、生成分组ID的方式实现：

# 标记起始行（b列为x）和结束行（b列为z）
df['start'] = df['b'] == 'x'
df['end'] = df['b'] == 'z'

# 生成分组ID：每次遇到start时递增，直到end之后
df['group_id'] = df['start'].cumsum()

# 计算每个分组的首尾索引范围
group_ranges = df.groupby('group_id').agg(
    first_idx=('start', 'idxmax'),
    last_idx=('end', 'idxmax')
)

# 提取每个分组的有效数据
result = []
for _, row in group_ranges.iterrows():
    result.append(df.loc[row['first_idx']:row['last_idx'], ['a', 'b']])

# 输出分组结果
for i, group in enumerate(result, 1):
    print(f"分组{i}:")
    print(group)
    print()

代码说明 #

1. 标记起止行：用布尔列start和end分别标记b列等于x和z的行。
2. 生成分组ID：通过cumsum()对start列累加，每遇到一个x就生成新的分组ID。
3. 确定分组边界：对每个分组ID，找到该组第一个起始行和最后一个结束行的索引，锁定分组范围。
4. 提取分组数据：根据边界索引提取对应子DataFrame，最终输出符合要求的分组结果。

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