python - Fillna 一次使用多种方法 - pandas

Question

我有一个如下所示的数据框

df = pd.DataFrame({'person_id': [101,101,101,101,202,202,202],
                        'start_date':['5/7/2013 09:27:00 AM','09/08/2013 11:21:00 AM','06/06/2014 08:00:00 AM','06/06/2014 05:00:00 AM','12/11/2011 10:00:00 AM','13/10/2012 12:00:00 AM','13/12/2012 11:45:00 AM'],
                        'end_date':['5/12/2013 09:27:00 AM',np.nan,'06/11/2014 08:00:00 AM',np.nan,'12/16/2011 10:00:00','10/18/2012 00:00:00',np.nan],
                        'type':['O','I','O','O','I','O','I']})
df.start_date = pd.to_datetime(df.start_date)
df['end_date'] = pd.to_datetime(df.end_date)

我想fillna()在end_date下基于以下两种方法的列

a) 如果NA在那个人的除最后一行以外的任何行中找到，fillna通过复制下一行的值

b) 如果 NA在那个人的最后一行中找到fillna在他的 start_date 上增加 10 天(因为没有下一行可供该人复制。所以，我们给出 10 天的随机值)

规则a和 b仅适用于 type=I 的人.

对于 type=O 的人, 只需复制 start_date 中的值即可填充.

这是我试过的。您可以看到我将同一行代码写了两次。

  df['end_date'] = np.where(df['type'].str.contains('I'),pd.DatetimeIndex(df['end_date'].bfill()),pd.DatetimeIndex(df.start_date.dt.date))
  df['end_date'] = np.where(df['type'].str.contains('I'),pd.DatetimeIndex(df['start_date'] + pd.DateOffset(10)),pd.DatetimeIndex(df.start_date.dt.date))

因为我必须将它应用于具有 1500 万行的大数据，所以有什么优雅有效的方式来编写它吗？

我希望我的输出如下所示

Answer 1

解决方案

s1 = df.groupby('person_id')['start_date'].shift(-1)
s1 = s1.fillna(df['start_date'] + pd.DateOffset(days=10))
s1 = df['end_date'].fillna(s1)

s2 = df['end_date'].fillna(df['start_date'])
df['end_date'] = np.where(df['type'].eq('I'), s1, s2)

说明

Group dataframe on person_id 和 shift start_date 列向上一个单位。

>>> df.groupby('person_id')['start_date'].shift(-1)

0   2013-09-08 11:21:00
1   2014-06-06 08:00:00
2   2014-06-06 05:00:00
3                   NaT
4   2012-10-13 00:00:00
5   2012-12-13 11:45:00
6                   NaT
Name: start_date, dtype: datetime64[ns]

在添加 10 天

的偏移量后，用 start_date 列中的值填充移位列中的 NaN 值

>>> s1.fillna(df['start_date'] + pd.DateOffset(days=10))

0   2013-09-08 11:21:00
1   2014-06-06 08:00:00
2   2014-06-06 05:00:00
3   2014-06-16 05:00:00
4   2012-10-13 00:00:00
5   2012-12-13 11:45:00
6   2012-12-23 11:45:00
Name: start_date, dtype: datetime64[ns]

现在用上述系列 s1 填充 end_date 列中的 NaN 值

>>>  df['end_date'].fillna(s1)

0   2013-05-12 09:27:00
1   2014-06-06 08:00:00
2   2014-06-11 08:00:00
3   2014-06-16 05:00:00
4   2011-12-16 10:00:00
5   2012-10-18 00:00:00
6   2012-12-23 11:45:00
Name: end_date, dtype: datetime64[ns]

类似地用 start_date 列中的值填充 end_date 列中的 NaN 值以创建一个系列 s2

>>> df['end_date'].fillna(df['start_date'])

0   2013-05-12 09:27:00
1   2013-09-08 11:21:00
2   2014-06-11 08:00:00
3   2014-06-06 05:00:00
4   2011-12-16 10:00:00
5   2012-10-18 00:00:00
6   2012-12-13 11:45:00
Name: end_date, dtype: datetime64[ns]

然后使用np.where根据类型为I<的条件从s1/s2中选择值 或 O

>>> df

   person_id          start_date            end_date type
0        101 2013-05-07 09:27:00 2013-05-12 09:27:00    O
1        101 2013-09-08 11:21:00 2014-06-06 08:00:00    I
2        101 2014-06-06 08:00:00 2014-06-11 08:00:00    O
3        101 2014-06-06 05:00:00 2014-06-06 05:00:00    O
4        202 2011-12-11 10:00:00 2011-12-16 10:00:00    I
5        202 2012-10-13 00:00:00 2012-10-18 00:00:00    O
6        202 2012-12-13 11:45:00 2012-12-23 11:45:00    I