高效计算pandas中的滚动时间差

Question

我在使用pandas这个工具，想要计算一个人在每个阶段花费的时间。为了让你更好地理解，我的数据集是这样的：

group       date    stage  
 A     2014-01-01   one   
 A     2014-01-03   one    
 A     2014-01-04   one    
 A     2014-01-05   two    
 B     2014-01-02  four    
 B     2014-01-06  five    
 B     2014-01-10  five    
 C     2014-01-03   two    
 C     2014-01-05   two    

我想要计算每个阶段的持续时间，以便得到：

 group       date    stage  dur
  A     2014-01-01   one    0
  A     2014-01-03   one    2
  A     2014-01-04   one    3
  A     2014-01-05   two    0
  B     2014-01-02  four    0
  B     2014-01-06  five    0
  B     2014-01-10  five    4
  C     2014-01-03   two    0
  C     2014-01-05   two    2

我下面使用的方法非常慢。有没有更快的方法呢？

df['stage_duration'] = df.groupby(['group', 'stage']).date.apply(lambda y: (y - y.iloc[0])).apply(lambda y:y / np.timedelta64(1, 'D')))

Answer 1

我觉得这里可以用 diff 来处理：

In [11]: df.groupby('stage')['date'].diff().fillna(0)
Out[11]:
0    0
1    2
2    0
3    0
4    0
5    4
dtype: float64

（假设这些阶段是连续的。）

如果你只是想从每组中减去第一个值，可以使用一个叫做 transform 的方法：

In [21]: df['date'] - df.groupby('stage')['date'].transform(lambda x: x.iloc[0])
Out[21]:
0    0
1    2
2    0
3    0
4    0
5    4
Name: date, dtype: int64

注意：这样做可能会快很多……

Answer 2

根据你的代码（你的 groupby/apply），看起来（尽管你的例子……但也许我误解了你的意思，那么安迪的做法可能是最好的选择）你正在处理一个 'date' 列，这个列的数据类型是 datetime64，而不是 integer。同时，看起来你想计算从某个 group/stage 的第一次观察开始的天数变化。我觉得这个例子数据更合适（如果我理解你的目标没错的话）：

>>> df

  group       date stage  dur
0     A 2014-01-01   one    0
1     A 2014-01-03   one    2
2     A 2014-01-04   one    3
3     A 2014-01-05   two    0
4     B 2014-01-02  four    0
5     B 2014-01-06  five    0
6     B 2014-01-10  five    4
7     C 2014-01-03   two    0
8     C 2014-01-05   two    2

考虑到你可以通过修改你的 apply 方法来提高速度（正如杰夫在他的评论中提到的），可以在 apply 之后以矢量化的方式通过 timedelta64 来进行除法运算（或者你也可以在 apply 中完成这一步）：

>>> df['dur'] = df.groupby(['group','stage']).date.apply(lambda x: x - x.iloc[0])
>>> df['dur'] /= np.timedelta64(1,'D')
>>> df

  group       date stage  dur
0     A 2014-01-01   one    0
1     A 2014-01-03   one    2
2     A 2014-01-04   one    3
3     A 2014-01-05   two    0
4     B 2014-01-02  four    0
5     B 2014-01-06  five    0
6     B 2014-01-10  five    4
7     C 2014-01-03   two    0
8     C 2014-01-05   two    2

但是，鉴于你的数据已经按 group、stage 和 date 排序，你也可以避免使用 groupby/apply。每个 ['group','stage'] 分组的第一个日期出现在组或阶段发生变化时。所以我觉得你可以做类似以下的操作：

>>> beg = (df.group != df.group.shift(1)) | (df.stage != df.stage.shift(1))
>>> df['dur'] = (df['date'] - df['date'].where(beg).ffill())/np.timedelta64(1,'D')
>>> df

  group       date stage  dur
0     A 2014-01-01   one    0
1     A 2014-01-03   one    2
2     A 2014-01-04   one    3
3     A 2014-01-05   two    0
4     B 2014-01-02  four    0
5     B 2014-01-06  five    0
6     B 2014-01-10  five    4
7     C 2014-01-03   two    0
8     C 2014-01-05   two    2

解释一下：注意 df['date'].where(beg) 创建了什么：

>>> beg = (df.group != df.group.shift(1)) | (df.stage != df.stage.shift(1))
>>> df['date'].where(beg)

0   2014-01-01
1          NaT
2          NaT
3   2014-01-05
4   2014-01-02
5   2014-01-06
6          NaT
7   2014-01-03
8          NaT

然后我用 ffill 填充这些值，并与 'date' 列进行差值计算。

编辑：正如安迪指出的，你也可以使用 transform：

>>> df['dur'] = df.date - df.groupby(['group','stage']).date.transform(lambda x: x.iloc[0])
>>> df['dur'] /= np.timedelta64(1,'D')

  group       date stage  dur
0     A 2014-01-01   one    0
1     A 2014-01-03   one    2
2     A 2014-01-04   one    3
3     A 2014-01-05   two    0
4     B 2014-01-02  four    0
5     B 2014-01-06  five    0
6     B 2014-01-10  five    4
7     C 2014-01-03   two    0
8     C 2014-01-05   two    2

速度：我用一个有 400,000 条观察数据的类似数据框测试了这两种方法的时间：

Apply 方法：

1 loops, best of 3: 18.3 s per loop

非 apply 方法：

1 loops, best of 3: 1.64 s per loop

所以我认为避免使用 apply 可以显著提高速度。