如果数据框中年份相邻，如何按组选择行？

Question

我有一个熊猫数据框，看起来像这样：

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

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

生成的数据帧应该如下所示

       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”，因为2015年和2017年之间存在“差距”

Answer 1

将date列转换为时间序列，然后使用^{}将a列上的year分组，并使用diff获得连续的差异，然后使用shift+ffill，最后使用abs+Series.eq创建布尔掩码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

Answer 2

为了清晰起见，这里有一个使用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]

输出：

       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