如何使用Pandas的groupby实现组与前一组的重叠分组？

原始数据 #

import pandas as pd

df = pd.DataFrame(
    {
        'a': list('xxxxxxxxxxyyyyyyyyy'),
        'b': list('1111222333112233444')
    }
)

需求说明 #

需要生成重叠分组，具体规则：

• 先按['a', 'b']分组，得到每个连续的同a同b数据块
• 将每个数据块与前一个同a值的相邻数据块合并，形成重叠分组
• 分组结果必须属于同一个a值，不能同时包含x和y

期望输出的重叠分组示例如下：

a  b
0   x  1
1   x  1
2   x  1
3   x  1
4   x  2
5   x  2
6   x  2

    a  b
4   x  2
5   x  2
6   x  2
7   x  3
8   x  3
9   x  3

    a  b
10  y  1
11  y  1
12  y  2
13  y  2

    a  b
12  y  2
13  y  2
14  y  3
15  y  3

    a  b
14  y  3
15  y  3
16  y  4
17  y  4
18  y  4

实现代码 #

1. 标记连续分组 #

先给每个['a', 'b']的连续块分配唯一ID，确保同一a值内的分组ID连续：

# 检测[a,b]列的变化，生成分组ID
df['group_id'] = (df[['a', 'b']] != df[['a', 'b']].shift()).any(axis=1).cumsum()

2. 生成重叠分组 #

按a值拆分处理，将每个a值下的相邻分组合并：

overlap_groups = []

# 遍历每个a值对应的子数据集
for a_val, sub_df in df.groupby('a'):
    # 获取当前a值下的所有分组ID
    ids = sub_df['group_id'].unique()
    # 遍历相邻的分组ID对
    for i in range(1, len(ids)):
        # 筛选出前一个分组和当前分组的所有行
        selected = sub_df[sub_df['group_id'].isin([ids[i-1], ids[i]])][['a', 'b']]
        overlap_groups.append(selected)

3. 输出结果 #

遍历overlap_groups即可得到所有重叠分组：

for idx, group in enumerate(overlap_groups):
    print(f"重叠分组 {idx+1}:")
    print(group)
    print()

运行上述代码后，输出结果将与需求示例完全一致。

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