比较两个日期列，然后根据差异在熊猫中设置一个标志

Question

我有一个带有两个datetime列的pandas数据框，如下所示，

d1  d2
0   2016-07-05  2016-08-15
1   2016-09-14  2016-10-11
2   2016-10-12  2016-11-08
3   2016-11-09  2016-12-06
4   2016-12-07  2017-01-03
5   2017-01-04  2017-02-28
6   2017-03-01  2017-03-28
7   2017-04-03  2017-04-30
8   2017-05-03  2017-05-30
9   2017-06-02  2017-06-29
10  2019-04-25  2019-05-22

我想根据条件创建名为flag的列

如果最后一个d2和d1分开超过1天，则将标志设置为最后一个标志+1。否则，使用相同的标志。对于第一行，标记始终为1。

我正在寻找的结果是

0     1
1     2
2     2
3     2
4     2
5     2
6     2
7     3
8     4
9     5
10    6

我正在使用很长的代码来实现这一目标，

(~test['d1'].sub(test['d2'].shift()).fillna(pd.Timedelta(days=0)).le(pd.Timedelta(days=1))).cumsum() + 1

我担心这行代码的效率。有什么更好的选择吗？

Answer 1

我会结合使用shift和Assign和cumsum方法。看下面的代码：

df = (df
     .assign(d1 = lambda d: pd.to_datetime(d.d1),
             d2 = lambda d: pd.to_datetime(d.d2),
             difference = lambda d: (d.d1-d.d2.shift()).dt.days.abs(),
             flag = lambda d: 1+(d.difference>1).cumsum()))