检查pandas数据帧中列中的几个条件的最有效方法是什么

Question

我正在处理一个包含三个相关列和270万行的pandas数据帧。结构为：

key VisitLink   dx_filter   time
0   1   ddcde14 1   100
1   2   abcde11 1   140
2   3   absdf12 1   50
3   4   ddcde14 0   125
4   5   ddcde14 1   140
data = [[1,'ddcde14',1,100],[2,'abcde11',1,140],[3,'absdf12',1,50],[4,'ddcde14',0,125],[5,'ddcde14',1,140]]
df_example = pd.DataFrame(data,columns = ['key','VisitLink','dx_filter','time'])

我需要3件事才是真的：
-VisitLink：两行之间的匹配
-dx_filter第一个事件为1
->strong>时间第二个事件发生在第一个事件的30天内示例：密钥1将生成密钥4作为匹配记录，因为它满足所有条件，但密钥4不会生成密钥5，因为它的dx_filter=0。

我进行了一次试验，我预测我的方法将需要120多个小时才能完成，我想知道是否有办法将其缩短到<10个小时，如果不可能的话。

def add_readmit_id(df):
df['readmit_id'] = np.nan
def set_id(row):
if row['dx_filter'] ==0:
return np.nan
else:
relevant_df = df.loc[df['VisitLink']==row['VisitLink']]
timeframe_df = relevant_df.loc[(relevant_df['time']>row['time'])&(relevant_df['time']<=row['time']+30)]
next_timeframe = timeframe_df['time'].min()
id_row = timeframe_df.loc[timeframe_df['time']==next_timeframe]
if not id_row.empty:
return id_row.iloc[0]['key']
else:
return np.nan
df['readmit_id'] = df.apply(set_id,axis=1)
return df
df_example = add_readmit_id(df_example)

请参阅上面我用来运行它的代码#最小可重复性。

Answer 1

以下是我使用groupby:的方法

groups = df.groupby('VisitLink')
s = groups['time'].diff(-1).le(30) & df['dx_filter']
df['shifted'] = np.where(s, groups['key'].shift(-1), np.nan)

输出：

key VisitLink  dx_filter  time  shifted
0    1   ddcde14          1   100      4.0
1    2   abcde11          1   140      NaN
2    3   absdf12          1    50      NaN
3    4   ddcde14          0   125      NaN
4    5   ddcde14          1   140      NaN