优雅地编写np.WHERE表示列中的不同值

问题描述

我有如下所示的数据帧

df = pd.DataFrame({'person_id': [101,101,101,101,202,202,202],
                   'person_type':['A','A','B','C','D','B','A'],
                   'login_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'],
                   'logout_date':[np.nan,'11/08/2013 11:21:00 AM',np.nan,'06/06/2014 05:00:00 AM',np.nan,'13/10/2012 12:00:00 AM',np.nan]})
df.login_date = pd.to_datetime(df.login_date)
df.logout_date = pd.to_datetime(df.logout_date)

我要将2个规则应用到logout_date列

规则1-如果人员类型为B、C、D、E且logout_date为NaN，则复制登录日期值

规则2-如果人员类型为A且logout_date为NaN，则将登录日期增加2天

我尝试了下面的

df['logout_date'] = np.where(((df['person_type'].isin(['B','C','D'])) & (df['logout_date'].isna())),df['login_date'].dt.date,df['logout_date'].dt.date)
df['logout_date'] = np.where(((df['person_type'].isin(['A'])) & (df['logout_date'].isna())),df['login_date'] + pd.DateOffset(days=2).dt.date,df['logout_date'].dt.date)

你可以看到它有多长。有没有其他更好的方法来写这篇文章？

我希望我的输出如下所示

person_id   person_type login_date           logout_date
101            A        2013-05-07 09:27:00  2013-05-09 09:27:00
101            A        2013-09-08 11:21:00  2013-11-08 11:21:00
101            B        2014-06-06 08:00:00  2014-06-06 08:00:00
101            C        2014-06-06 05:00:00  2014-06-06 05:00:00
202            D        2011-12-11 10:00:00  2011-12-11 10:00:00
202            B        2012-10-13 00:00:00  2012-10-13 12:00:00
202            A        2012-12-13 11:45:00  2012-12-15 11:45:00

---

解决方案

对注释中提到的中间变量使用numpy.select：

s = df['person_type'].fillna('missing value')
m1 = s.isin(['B','C','D', 'missing value'])
m2 = s.isin(['A','missing value'])

df['logout_date'] = np.select([m1, m2],
                              [df['login_date'], df['login_date'] + pd.DateOffset(days=2)],
                               default=df['logout_date'])

或重写您的解决方案：

m1 = df['person_type'].isin(['B','C','D'])
m2 = df['person_type'].isin(['A'])
m3 = df['logout_date'].isna()

df['logout_date'] = np.select([m1 & m3, m2 & m3],
                              [df['login_date'], df['login_date'] + pd.DateOffset(days=2)],
                               default=df['logout_date'])

---

df['logout_date'] = np.select([m1 & m3, m2 & m3],
                              [df['login_date'].dt.date, 
                               (df['login_date'] + pd.DateOffset(days=2)).dt.date],
                               default=df['logout_date'].dt.date)