Pandas如何提取DataFrame中flag列连续为1且大小大于2的分组

示例输入 #

import pandas as pd
df2 = pd.DataFrame({
    'A':[1,20,40,45,56,1,20,40,45,56],
    'flag':[3,2,4,1,1,3,3,1,1,1]
})
print(df2)

原始数据输出：

A  flag
0   1     3
1  20     2
2  40     4
3  45     1
4  56     1
5   1     3
6  20     3
7  40     1
8  45     1
9  56     1

预期输出 #

A  flag
7  40     1
8  45     1
9  56     1

实现代码 #

核心逻辑是先给连续相同的flag值分配独立分组ID，再按分组规则做筛选：

# 为连续相同的flag值生成分组标识
df2["group_id"] = (df2["flag"] != df2["flag"].shift()).cumsum()

# 筛选flag=1、且分组行数大于2的记录，删除临时生成的分组ID列
res = df2[
    (df2["flag"] == 1)
    & (df2.groupby("group_id")["group_id"].transform("count") > 2)
].drop(columns="group_id")

print(res)

运行代码后输出结果和预期完全一致。

注：索引3、4位置的连续1分组长度为2，不满足长度大于2的筛选条件，因此不会被保留。

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