如何在Pandas DataFrame中找出连续0和1的重叠分组？

有如下简单的DataFrame，包含二进制列status，以及可用于聚合的数值列column1和column2：

>>> df
   status  column1  column2
0       1       10        5
1       0       20        4
2       0       30        3
3       1       40        2
4       0       50        1
5       0       60        5
6       0       70        4
7       1       80        3
8       1       90        2
9       1      100        1
10      0      110        5

需要基于status列创建重叠分组：分组由1后跟0、0后跟1的状态切换点触发，形成的分组范围为（从索引0开始）：索引0至2、索引2至3、索引3至6、索引6至9、索引9至10。

之后对分组后的数值列执行聚合操作：对column1求和，对column2求中位数，预期输出如下：

column1  column2
0     60.0      4.0
1     70.0      2.5
2    220.0      3.0
3    340.0      2.5
4    210.0      3.0

此前尝试过df['status'].diff().ne(x).cumsum()的多种变体，但无法处理重叠分组的场景。

要实现这种重叠分组的聚合，可以通过定位状态切换点、生成分组索引对、逐个计算聚合值的步骤完成：

步骤1：定位状态切换点 #

找出status列中值发生变化的索引位置，这些位置就是分组的分界点：

# 获取status列发生变化的索引（跳过第一个元素，因为shift后第一个值为NaN）
switch_points = df['status'].diff().ne(0).iloc[1:].index.tolist()
# 补充起始索引和最后一行索引，形成完整的分界点列表
all_points = [0] + switch_points + [df.index[-1]]

步骤2：生成重叠分组的起止索引对 #

相邻分界点之间形成一个分组，且相邻分组共享分界点（实现重叠）：

groups = [(all_points[i], all_points[i+1]) for i in range(len(all_points)-1)]

步骤3：计算分组聚合值 #

遍历每个分组，分别计算column1的和与column2的中位数，最后整理成结果DataFrame：

import pandas as pd

results = []
for start, end in groups:
    group_data = df.loc[start:end]
    sum_col1 = group_data['column1'].sum()
    median_col2 = group_data['column2'].median()
    results.append({'column1': sum_col1, 'column2': median_col2})

result_df = pd.DataFrame(results)
print(result_df)

运行代码后输出结果与预期一致：

column1  column2
0     60.0      4.0
1     70.0      2.5
2    220.0      3.0
3    340.0      2.5
4    210.0      3.0

简化实现 #

可以将步骤合并，用更简洁的代码完成：

# 获取所有状态切换点
switch_idx = df['status'].ne(df['status'].shift()).iloc[1:].index.tolist()
all_idx = [0] + switch_idx + [df.index[-1]]

# 直接生成聚合结果
result_df = pd.DataFrame([
    {
        'column1': df.loc[s:e, 'column1'].sum(),
        'column2': df.loc[s:e, 'column2'].median()
    } for s, e in zip(all_idx[:-1], all_idx[1:])
])

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