如果 pandas dataframe 中的年份彼此相邻，如何按组划分 select 行？

Question

我有一个 pandas dataframe 看起来像这样：

pd.DataFrame({'a':['cust1', 'cust1', 'cust2', 'cust3', 'cust3', 'cust4', 'cust4'],
                   'date':[date(2019, 1, 20), date(2020, 6, 15), date(2017, 1, 10), date(2015, 6, 12), date(2017, 12, 15), date(2018, 12, 10), date(2017, 1, 5)]
                   })

       a        date
0  cust1  2019-01-20
1  cust1  2020-06-15
2  cust2  2017-01-10
3  cust3  2015-06-12
4  cust3  2017-12-15
5  cust4  2018-12-10
6  cust4  2017-01-05

我想 select 如果“日期”列中的年份是“背靠背”或彼此相邻（即年份之间没有间隔）并且每组至少有 2 个不同的日期，则在“a”列中按组对所有行在一个'。

生成的 dataframe 应如下所示。

       a        date
0  cust1  2019-01-20
1  cust1  2020-06-15
5  cust4  2018-12-10
6  cust4  2017-01-05

解释：没有选择第 2 行，因为组 'cust2' 只有 1 年/日期。 组“cust3”不是 select，因为 2015 年和 2017 年之间存在“差距”。

Answer 1

这是一个使用groupby和shift的解决方案，为了清楚起见，分几个步骤：

df["year"] = pd.to_datetime(df.date).dt.year
df["adjacent"] = (df.groupby("a")["year"].diff().abs() == 1) | (df.groupby("a")["year"].diff(-1).abs() == 1)
res = df[df.adjacent]

Output：

       a       date  year    adjacent
0  cust1 2019-01-20  2019    True
1  cust1 2020-06-15  2020    True
5  cust4 2018-12-10  2018    True
6  cust4 2017-01-05  2017    True

Answer 2

将date列转换为 pandas 日期时间序列，然后使用Series.groupby将a列上的year分组并使用diff获取连续差异，然后shift + ffill ，最后使用abs + Series.eq创建一个 boolean 掩码m ，使用这个掩码过滤行：

df['date'] = pd.to_datetime(df['date'])
m = df['date'].dt.year.groupby(df['a']).apply(lambda x: x.diff().shift(-1).ffill()).abs().eq(1)
df1 = df[m]

结果：

print(df1)
       a       date
0  cust1 2019-01-20
1  cust1 2020-06-15
5  cust4 2018-12-10
6  cust4 2017-01-05