如何基于is_hot列实现可中断的高温日连续天数（Streak）统计

问题背景 #

现有如下DataFrame，需要统计高温日的连续天数（streak）：要求当is_hot值为0时，连续天数中断并重新计数，同时将NaN值填充为0。

创建DataFrame的代码 #

import pandas as pd
import numpy as np

dates = pd.date_range(start ='1-1-2018', end ='1-10-2018', freq ='1D')
np.random.seed(42)
temp = np.random.randint(60, 80, size=10)
df = pd.DataFrame({'dates': dates, 'temp':temp})
df["is_hot"] = np.where(df["temp"] > 70, 1, 0)

初始DataFrame #

原代码问题 #

使用以下代码统计时，会出现非连续高温日的streak累计不中断，且非高温日为NaN的情况：

df['streak'] = df.loc[df['is_hot'].eq(1)].groupby(df['is_hot'])['is_hot'].cumsum()

原代码输出结果 #

期望结果：2018-01-07的streak值为1，所有NaN填充为0。

解决方案 #

核心思路：先标记连续相同状态的分组，再在每个分组内累计计数，非高温日分组自动累计为0。

代码实现 #

# 标记连续状态的分组：当is_hot状态变化时，分组ID递增
df['group'] = (df['is_hot'] != df['is_hot'].shift(1)).cumsum()

# 在每个分组内对is_hot累加，非高温日分组累加结果为0
df['streak'] = df.groupby('group')['is_hot'].cumsum()

# 可选：删除临时的group列
df = df.drop('group', axis=1)

最终结果 #

代码说明 #

• (df['is_hot'] != df['is_hot'].shift(1)).cumsum()：通过比较当前行与上一行的is_hot值，生成分组ID，每当状态从0变1或1变0时，分组ID递增，实现连续相同状态的分组。
• df.groupby('group')['is_hot'].cumsum()：在每个分组内对is_hot累加，非高温日分组的is_hot全为0，累加结果自然为0；高温日分组则从1开始连续计数，中断后新分组重新从1开始。

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