Pandas实现数值区间groupby分组并提取起止值方法

请问是否可参照下述示例在Pandas中对整数类型数值区间执行groupby分组操作？若无法直接实现，应通过何种方式得到预期输出结果？

测试数据集构建代码如下：

df = pd.DataFrame(
    {"price": [9, 8, 9, 10, 11, 6, 7, 8, 9, 9, 9, 9, 10, 11, 5]},
    index=pd.date_range("19/3/2020", periods=15, freq="H"),
)

df["higher"] = np.where(df.price > df.price.shift(), 1, 0)
df["higher_count"] = df["higher"] * (
    df["higher"].groupby((df["higher"] != df["higher"].shift()).cumsum()).cumcount() + 1
)
df = df.drop("higher", axis=1)

首个分组已高亮标注的Dataframe示例：

现有基于itertools的分组提取代码如下：

from operator import itemgetter
from itertools import groupby

data = df["higher_count"]
for key, group in groupby(enumerate(data), lambda i: i[0] - i[1]):         
    group = list(map(itemgetter(1), group))
    if len(group) > 1:
        print(f"{key}:{group}")

# 执行输出
# 1:[0, 1, 2, 3]
# 5:[0, 1, 2, 3]
# 11:[0, 1, 2]

需要为每个符合条件的分组生成如下字段：

• start date（起始日期）
• price at start date（起始日期对应价格）
• end date（结束日期）
• price at end date（结束日期对应价格）

其中key为1的分组对应的输出示例如下：

完全可以用Pandas原生语法实现，不需要借助itertools做逐行循环，性能更好，逻辑也更清晰。核心思路是先给每一段连续上涨的区间生成唯一分组ID，再按组聚合提取首尾的时间和价格即可。

完整代码 #

import pandas as pd
import numpy as np

# 构建测试数据
df = pd.DataFrame(
    {"price": [9, 8, 9, 10, 11, 6, 7, 8, 9, 9, 9, 9, 10, 11, 5]},
    index=pd.date_range("2020-03-19", periods=15, freq="H"),
)

# 1. 标记当前行相比上一行是否上涨
df["is_rise"] = df["price"] > df["price"].shift()
# 2. 状态变化时分组ID+1，同一段连续状态共享同一个ID
df["group_id"] = (df["is_rise"] != df["is_rise"].shift()).cumsum()

result = []
# 3. 按分组遍历，提取符合要求的连续上涨段信息
for gid, sub_df in df.groupby("group_id"):
    rise_seg = sub_df[sub_df["is_rise"]]
    # 过滤掉长度不足的非连续上涨段，和原示例len(group)>1的筛选逻辑对齐
    if len(rise_seg) < 1:
        continue
    # 连续上涨段的起始点是该组第一行（上涨前的基准点），结束点是上涨段最后一行
    start = sub_df.iloc[0]
    end = rise_seg.iloc[-1]
    result.append({
        "start date": start.name,
        "price at start date": start["price"],
        "end date": end.name,
        "price at end date": end["price"]
    })

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

运行结果 #

start date  price at start date            end date  price at end date
0 2020-03-19 01:00:00                    8 2020-03-19 04:00:00                 11
1 2020-03-19 05:00:00                    6 2020-03-19 08:00:00                  9
2 2020-03-19 11:00:00                    9 2020-03-19 13:00:00                 11

第一行结果就是key=1分组的输出，和示例截图完全匹配。

逻辑说明 #

这种相邻状态变化生成分组ID的写法是Pandas处理连续区间分组的通用方案：通过对比当前行和上一行的状态，在状态发生切换时给分组ID累加1，就能自动给每一段连续不变的区间分配独立ID，后续不管做聚合、统计都非常方便，比循环遍历的执行效率高几个量级，适配百万级以上的数据集。

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