如何按director_id和match_id筛选含至少连续两年数据的观测值？

我有一个包含director_id、match_id、calyear字段的数据集，需要仅保留按director_id和match_id分组后，组内存在至少连续两年数据的所有观测值。之前尝试的方法步骤繁琐还存在问题，希望得到更简洁的解决方案。

现有数据集 #

期望结果 #

原方法存在的问题 #

• 操作流程繁琐，需额外创建辅助列
• 分组逻辑有误，导致组内首行被设为NaN，无法正确覆盖所有符合条件的观测

简洁实现思路 #

核心逻辑：先按分组对年份排序，检查组内是否存在当前年份与相邻年份连续的情况，标记出所有符合条件的分组，最后保留这些分组的全部观测值。

具体代码（Pandas） #

import pandas as pd

# 初始化示例数据集
df = pd.DataFrame([
    [282, 1111, 2006], [282, 1111, 2007],
    [356, 2222, 2005], [356, 2222, 2007],
    [600, 3333, 2010], [600, 3333, 2011], [600, 3333, 2012],
    [600, 3355, 2013], [600, 3355, 2015], [600, 3355, 2016],
    [753, 4444, 2005], [753, 4444, 2008], [753, 4444, 2009]
], columns=['director_id', 'match_id', 'calyear'])

# 1. 按分组+年份排序，保证后续判断的准确性
df_sorted = df.sort_values(by=['director_id', 'match_id', 'calyear'])

# 2. 标记分组是否存在连续年份
df_sorted['has_consecutive'] = df_sorted.groupby(['director_id', 'match_id'])['calyear'].transform(
    lambda x: (x.diff() == 1).any() | (x.shift(-1) - x == 1).any()
)

# 3. 筛选符合条件的观测，删除辅助列
result = df_sorted[df_sorted['has_consecutive']].drop(columns='has_consecutive').reset_index(drop=True)

print(result)

代码解释 #

• 排序：确保每个分组内的年份是升序排列，避免因年份乱序导致连续判断出错。
• 标记连续分组：
  • x.diff() == 1：检查当前年份是否比上一年大1（捕获前序连续的情况）
  • x.shift(-1) - x == 1：检查当前年份是否比下一年小1（捕获后序连续的情况，覆盖连续组的最后一行）
  • any()判断整个分组是否存在至少一组连续年份，返回布尔值标记每个观测
• 筛选结果：直接过滤出标记为True的观测，删除辅助列后得到最终结果。

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