如何在Pandas中按连续stage值统计非None数量并生成num_issues列？

现有如下Pandas DataFrame：

temp   stage   issue_datetime
20      1      2022/11/30 19:20
21      1      2022/11/30 19:21
20      1      None
25      1      2022/11/30 20:10
30      2      None
22      2      2022/12/01 10:00
22      2      2022/12/01 10:01
31      3      2022/12/02 11:00
32      3      2022/12/02 11:01
19      1      None
20      1      None

希望得到的结果：

temp   stage   num_issues
20      1      3
21      1      3
20      1      3
25      1      3
30      2      2
22      2      2
22      2      2
31      3      2
32      3      2
19      1      0
20      1      0

需求：按连续的stage值分组，统计每组中issue_datetime列的非None数量，生成名为num_issues的新列。

核心逻辑是先识别出连续的stage区块（而非全局按stage值分组），再对每个区块统计非空数量，最后映射回原DataFrame。

步骤1：生成连续分组标识 #

通过对比当前行与上一行的stage值，为每段连续相同的stage分配唯一分组ID：

import pandas as pd

# 假设原始数据存储在df中
df['group_id'] = (df['stage'] != df['stage'].shift()).cumsum()

步骤2：统计每个分组的非空数量 #

对每个分组ID，计算issue_datetime列非空值的个数：

counts = df.groupby('group_id')['issue_datetime'].apply(lambda x: x.notna().sum())

步骤3：映射统计结果到原数据 #

将统计好的数量合并回原DataFrame，并重命名为目标列名：

df['num_issues'] = df['group_id'].map(counts)
# 不需要临时分组列的话可删除
df = df.drop('group_id', axis=1)

完整可运行代码 #

import pandas as pd

# 构造原始数据
data = {
    'temp': [20,21,20,25,30,22,22,31,32,19,20],
    'stage': [1,1,1,1,2,2,2,3,3,1,1],
    'issue_datetime': ['2022/11/30 19:20', '2022/11/30 19:21', None, '2022/11/30 20:10', None, '2022/12/01 10:00', '2022/12/01 10:01', '2022/12/02 11:00', '2022/12/02 11:01', None, None]
}
df = pd.DataFrame(data)

# 生成连续分组标识
df['group_id'] = (df['stage'] != df['stage'].shift()).cumsum()

# 统计每组非空数量
counts = df.groupby('group_id')['issue_datetime'].apply(lambda x: x.notna().sum())

# 映射到原数据
df['num_issues'] = df['group_id'].map(counts)

# 清理临时列
df = df.drop('group_id', axis=1)

print(df)

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