python - 检查 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'])

我需要三件事是真实的:
- VisitLink:两行之间的匹配
- dx_filter:第一个事件为1
- 时间:第二个事件发生在第一个事件后的 30 天内示例: key 1 将生成 key 4 作为匹配记录，因为它满足所有条件，但 key 4 不会生成 key 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)

请参阅上面我用来运行它的代码#minimum reproducible。

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