python - 如果熊猫数据框中的年份彼此相邻，如何按组选择行？

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

如果“日期”列中的年份是“背靠背”或彼此相邻（即年份之间没有间隔）并且每组至少有 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 年之间存在“差距”。

标签： pythonpandasdataframepandas-groupby

Answer 1

这是一个使用groupbyand的解决方案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