如何按id筛选连续value=1的5行以构建平衡面板数据？

需求说明 #

从给定的DataFrame中，为每个id筛选出**连续5个program_year且value == 1**的记录，后续用于构建平衡面板数据。

原始数据 #

id  program_year  value
1     2007         1
1     2008         1
1     2009         1
1     2010         1
1     2011         1
1     2012         1
1     2013         0
2     2007         0
2     2008         1
2     2009         1
2     2010         1
2     2011         1  
2     2012         1
2     2013         1
3     2007         1
3     2008         0
3     2009         1
3     2010         1
3     2011         1
3     2012         1
3     2013         1

预期输出 #

id  program_year  value
1     2007         1
1     2008         1
1     2009         1
1     2010         1
1     2011         1
2     2008         1
2     2009         1
2     2010         1
2     2011         1  
2     2012         1
3     2009         1
3     2010         1
3     2011         1
3     2012         1
3     2013         1

解决方案（基于Pandas） #

通过分组标记连续年份区间，筛选出长度≥5的区间并提取对应记录：

import pandas as pd

# 构造原始数据框
data = {
    'id': [1]*7 + [2]*7 + [3]*7,
    'program_year': [2007,2008,2009,2010,2011,2012,2013]*3,
    'value': [1,1,1,1,1,1,0, 0,1,1,1,1,1,1, 1,0,1,1,1,1,1]
}
df = pd.DataFrame(data)

# 1. 过滤仅保留value=1的记录
df_filtered = df[df['value'] == 1].copy()

# 2. 按id分组，标记连续年份的分组
# 计算当前年份与前一年的差值，差值≠1则为新的连续组起点
df_filtered['year_diff'] = df_filtered.groupby('id')['program_year'].diff().fillna(1)
df_filtered['continuous_group'] = (df_filtered['year_diff'] != 1).groupby(df_filtered['id']).cumsum()

# 3. 筛选出长度≥5的连续组，每个id取该组的前5行（匹配预期输出）
# 先过滤符合长度要求的组
valid_groups = df_filtered.groupby(['id', 'continuous_group']).filter(lambda x: len(x) >=5)
# 每个id提取对应连续组的前5条记录
final_result = valid_groups.groupby('id').apply(lambda x: x.head(5)).reset_index(drop=True)

# 查看结果
print(final_result)

代码说明 #

• 第一步过滤掉value≠1的行，聚焦目标数据；
• 第二步通过diff()计算年份差，用cumsum()给每个连续年份区间分配唯一标识；
• 第三步筛选出长度≥5的连续区间，再按id提取每个区间的前5条记录，正好匹配预期输出的连续5年要求。

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